use std::collections::BTreeMap;

use crate::{
    ast::types::{
        ArrayDimension, BinaryOp, BlockStmt, BreakStmt, CompileUnit, ContinueStmt, Expr, ExprValue, FuncDeclStmt,
        GlobalDeclStmt, IfElseBranch, IfStmt, ReturnStmt, Statement, VarDeclStmt, WhileStmt,
    },
    diagnostic::{span::Span, Diagnositics},
    sema::{
        err::SemaError,
        hir::{
            HirBlockStmt, HirBreakStmt, HirCompileUnit, HirContinueStmt, HirExpr, HirExprValue,
            HirFuncDeclStmt, HirGlobalDeclStmt, HirIfElseBranch, HirIfStmt, HirParam,
            HirReturnStmt, HirStatement, HirVarDeclStmt, HirVarDeclStmtValue, HirWhileStmt,
        },
        symbol::{FunctionId, FunctionSig, SymbolId, SymbolKind, SymbolTable},
        types::SemaType,
    },
};

pub struct Analyzer {
    symbols: SymbolTable,
    function_map: BTreeMap<String, FunctionId>,
    functions: Vec<FunctionSig>,
    current_func_return_type: Option<SemaType>,
    diagnostic: Diagnositics,
    loop_depth: usize,
}

impl Analyzer {
    pub fn new() -> Self {
        let mut analyzer = Self {
            symbols: SymbolTable::new(),
            function_map: BTreeMap::new(),
            functions: vec![],
            current_func_return_type: None,
            diagnostic: Diagnositics::new(),
            loop_depth: 0,
        };
        analyzer.declare_builtin_func("putint", vec![SemaType::I32], SemaType::Void);
        analyzer.declare_builtin_func("putch", vec![SemaType::I32], SemaType::Void);
        analyzer.declare_builtin_func("getint", vec![], SemaType::I32);
        analyzer
    }

    pub fn analyze(&mut self, compile_unit: CompileUnit) -> HirCompileUnit {
        self.analyze_compile_unit(compile_unit)
    }

    pub fn get_diagnostics(&self) -> &Diagnositics {
        &self.diagnostic
    }

    pub fn get_symbol_type(&self, symbol: SymbolId) -> SemaType {
        self.symbols.get_symbol(symbol).ty.clone()
    }

    pub fn get_symbol_kind(&self, symbol: SymbolId) -> SymbolKind {
        self.symbols.get_symbol(symbol).kind
    }

    pub fn get_function_sig(&self, function: FunctionId) -> &FunctionSig {
        &self.functions[function.0]
    }

    fn declare_builtin_func(&mut self, name: &str, parameter_types: Vec<SemaType>, return_type: SemaType) {
        let id = FunctionId(self.functions.len());
        let sig = FunctionSig {
            id,
            name: name.to_string(),
            return_type,
            parameter_types,
        };
        self.function_map.insert(name.to_string(), id);
        self.functions.push(sig);
    }

    fn add_error(&mut self, error: SemaError, span: Span) {
        self.diagnostic.add_from_error(error, span);
    }

    fn analyze_compile_unit(&mut self, compile_unit: CompileUnit) -> HirCompileUnit {
        let mut global_decls = vec![];
        for decl in compile_unit.global_decls {
            match decl {
                GlobalDeclStmt::VarDecl(var_decl) => {
                    global_decls.push(HirGlobalDeclStmt::VarDecl(self.analyze_var_decl(var_decl, SymbolKind::Global)));
                }
                GlobalDeclStmt::FuncDecl(func_decl) => {
                    if let Some(func_decl) = self.analyze_func_decl(func_decl) {
                        global_decls.push(HirGlobalDeclStmt::FuncDecl(func_decl));
                    }
                }
            }
        }
        HirCompileUnit { global_decls }
    }

    fn analyze_var_decl(&mut self, var_decl: VarDeclStmt, kind: SymbolKind) -> HirVarDeclStmt {
        let base_type: SemaType = var_decl.data_type.into();
        let mut values = vec![];
        for value in var_decl.values {
            let data_type = self.build_var_type(base_type.clone(), &value.dimensions, false);
            match self.symbols.declare_variable(&value.name, kind, data_type) {
                Ok(symbol) => values.push(HirVarDeclStmtValue {
                    symbol,
                    name_span: value.name_span,
                }),
                Err(e) => self.add_error(e, value.name_span),
            }
        }
        HirVarDeclStmt {
            values,
            data_type: base_type,
            type_span: var_decl.type_span,
        }
    }

    fn analyze_func_decl(&mut self, func_decl: FuncDeclStmt) -> Option<HirFuncDeclStmt> {
        if self.function_map.contains_key(&func_decl.name) {
            self.add_error(SemaError::FunctionHasBeenDefined(func_decl.name.clone()), func_decl.name_span);
            return None;
        }

        let function_id = FunctionId(self.functions.len());
        let mut parameter_types = vec![];
        for param in &func_decl.params {
            parameter_types.push(self.build_var_type(param.param_type.into(), &param.dimensions, true));
        }
        let return_type: SemaType = func_decl.return_type.into();
        let sig = FunctionSig {
            id: function_id,
            name: func_decl.name.clone(),
            return_type: return_type.clone(),
            parameter_types,
        };
        self.function_map.insert(func_decl.name.clone(), function_id);
        self.functions.push(sig.clone());

        self.current_func_return_type = Some(return_type);
        self.symbols.enter_scope();

        let mut params = vec![];
        for param in func_decl.params {
            let param_type = self.build_var_type(param.param_type.into(), &param.dimensions, true);
            match self.symbols.declare_variable(&param.name, SymbolKind::Param, param_type.clone()) {
                Ok(symbol) => params.push(HirParam {
                    symbol,
                    param_type,
                    name_span: param.name_span,
                    type_span: param.type_span,
                }),
                Err(e) => self.add_error(e, param.name_span),
            }
        }

        let body = self.analyze_block_stmt(func_decl.body);
        self.symbols.exit_scope();
        self.current_func_return_type = None;

        Some(HirFuncDeclStmt {
            sig,
            params,
            body,
            ret_type_span: func_decl.ret_type_span,
            name_span: func_decl.name_span,
        })
    }

    fn build_var_type(&mut self, base_type: SemaType, dimensions: &[ArrayDimension], is_param: bool) -> SemaType {
        if dimensions.is_empty() {
            return base_type;
        }
        let mut dims = vec![];
        for (i, dimension) in dimensions.iter().enumerate() {
            match &dimension.value {
                Some(expr) => {
                    if let ExprValue::IntLit(value) = expr.value {
                        if value > 0 {
                            dims.push(value as usize);
                        } else {
                            self.add_error(SemaError::InvalidArrayDimension, dimension.span);
                        }
                    } else {
                        self.add_error(SemaError::InvalidArrayDimension, dimension.span);
                    }
                }
                None if is_param && i == 0 => {}
                None => self.add_error(SemaError::InvalidArrayDimension, dimension.span),
            }
        }
        if is_param {
            let pointee_type = if dims.is_empty() {
                base_type
            } else {
                SemaType::Array(Box::new(base_type), dims)
            };
            SemaType::Ptr(Box::new(pointee_type))
        } else {
            SemaType::Array(Box::new(base_type), dims)
        }
    }

    fn analyze_block_stmt(&mut self, block_stmt: BlockStmt) -> HirBlockStmt {
        let mut statements = vec![];
        for stmt in block_stmt.statements {
            if let Some(stmt) = self.analyze_statement(stmt) {
                statements.push(stmt);
            }
        }
        HirBlockStmt { statements }
    }

    fn analyze_statement(&mut self, stmt: Statement) -> Option<HirStatement> {
        match stmt {
            Statement::Return(stmt) => Some(HirStatement::Return(self.analyze_return_stmt(stmt))),
            Statement::If(stmt) => Some(HirStatement::If(self.analyze_if_stmt(stmt))),
            Statement::While(stmt) => Some(HirStatement::While(self.analyze_while_stmt(stmt))),
            Statement::Break(stmt) => Some(HirStatement::Break(self.analyze_break_stmt(stmt))),
            Statement::Continue(stmt) => Some(HirStatement::Continue(self.analyze_continue_stmt(stmt))),
            Statement::Block(stmt) => {
                self.symbols.enter_scope();
                let block = self.analyze_block_stmt(stmt);
                self.symbols.exit_scope();
                Some(HirStatement::Block(block))
            }
            Statement::Expr(expr) => self.analyze_expr(expr).map(HirStatement::Expr),
            Statement::VarDecl(var_decl) => Some(HirStatement::VarDecl(self.analyze_var_decl(var_decl, SymbolKind::Local))),
        }
    }

    fn analyze_return_stmt(&mut self, return_stmt: ReturnStmt) -> HirReturnStmt {
        let expected_ty = self.current_func_return_type.clone().unwrap();
        let value = match return_stmt.value {
            Some(expr) => {
                if expected_ty == SemaType::Void {
                    self.add_error(SemaError::ReturnExpressionOnVoidFunction, return_stmt.span);
                    None
                } else {
                    match self.analyze_expr(expr) {
                        Some(expr) => {
                            if expr.ty == SemaType::Void {
                                self.add_error(SemaError::InvalidOperand(SemaType::Void), return_stmt.span);
                            } else if !self.type_matches(&expected_ty, &expr.ty) {
                                self.add_error(SemaError::TypeMismatch(expected_ty.clone(), expr.ty.clone()), return_stmt.span);
                            }
                            Some(expr)
                        }
                        None => None,
                    }
                }
            }
            None => {
                if expected_ty != SemaType::Void {
                    self.add_error(SemaError::TypeMismatch(expected_ty, SemaType::Void), return_stmt.span);
                }
                None
            }
        };
        HirReturnStmt {
            value,
            span: return_stmt.span,
        }
    }

    fn analyze_if_stmt(&mut self, if_stmt: IfStmt) -> HirIfStmt {
        let condition = self.analyze_condition_expr(if_stmt.condition);
        let then_branch = self.analyze_block_stmt(if_stmt.then_branch);
        let mut ifelse_branch = vec![];
        for branch in if_stmt.ifelse_branch {
            let IfElseBranch { condition, then_branch } = branch;
            ifelse_branch.push(HirIfElseBranch {
                condition: self.analyze_condition_expr(condition),
                then_branch: self.analyze_block_stmt(then_branch),
            });
        }
        let else_branch = if_stmt.else_branch.map(|block| self.analyze_block_stmt(block));
        HirIfStmt {
            condition,
            then_branch,
            ifelse_branch,
            else_branch,
        }
    }

    fn analyze_while_stmt(&mut self, while_stmt: WhileStmt) -> HirWhileStmt {
        let condition = self.analyze_condition_expr(while_stmt.condition);
        self.loop_depth += 1;
        let body = self.analyze_block_stmt(while_stmt.body);
        self.loop_depth -= 1;
        HirWhileStmt {
            condition,
            body,
        }
    }

    fn analyze_break_stmt(&mut self, stmt: BreakStmt) -> HirBreakStmt {
        if self.loop_depth == 0 {
            self.add_error(SemaError::BreakOutsideLoop, stmt.span);
        }
        HirBreakStmt { span: stmt.span }
    }

    fn analyze_continue_stmt(&mut self, stmt: ContinueStmt) -> HirContinueStmt {
        if self.loop_depth == 0 {
            self.add_error(SemaError::ContinueOutsideLoop, stmt.span);
        }
        HirContinueStmt { span: stmt.span }
    }

    fn analyze_condition_expr(&mut self, expr: Expr) -> HirExpr {
        let span = expr.span;
        match self.analyze_expr(expr) {
            Some(expr) => {
                if expr.ty == SemaType::Void {
                    self.add_error(SemaError::InvalidOperand(SemaType::Void), span);
                }
                expr
            }
            None => HirExpr {
                value: HirExprValue::IntLit(0),
                ty: SemaType::I32,
                span,
            },
        }
    }

    fn analyze_expr(&mut self, expr: Expr) -> Option<HirExpr> {
        let span = expr.span;
        match expr.value {
            ExprValue::IntLit(value) => Some(HirExpr {
                value: HirExprValue::IntLit(value),
                ty: SemaType::I32,
                span,
            }),
            ExprValue::Var(name) => {
                let symbol = match self.symbols.get_variable(&name) {
                    Some(symbol) => symbol,
                    None => {
                        self.add_error(SemaError::VariableNotFound(name), span);
                        return None;
                    }
                };
                Some(HirExpr {
                    value: HirExprValue::Var(symbol),
                    ty: self.symbols.get_symbol(symbol).ty.clone(),
                    span,
                })
            }
            ExprValue::Assign { lvalue, rvalue } => self.analyze_assign_expr(*lvalue, *rvalue, span),
            ExprValue::ArrayAccess { array, index } => self.analyze_array_access_expr(*array, *index, span),
            ExprValue::UnaryOp { op, operand } => {
                let operand_span = operand.span;
                let operand = self.analyze_expr(*operand)?;
                if operand.ty == SemaType::Void {
                    self.add_error(SemaError::InvalidOperand(SemaType::Void), operand_span);
                    return None;
                }
                let ty = match op {
                    crate::ast::types::UnaryOp::Not => SemaType::I1,
                    _ => operand.ty.clone(),
                };
                Some(HirExpr {
                    value: HirExprValue::UnaryOp {
                        op,
                        operand: Box::new(operand),
                    },
                    ty,
                    span,
                })
            }
            ExprValue::BinaryOp { lhs, op, rhs } => self.analyze_binary_expr(*lhs, op, *rhs, span),
            ExprValue::FuncCall(func_name, args) => self.analyze_func_call_expr(func_name, args, span),
        }
    }

    fn analyze_assign_expr(&mut self, lvalue: Expr, rvalue: Expr, span: Span) -> Option<HirExpr> {
        if !matches!(lvalue.value, ExprValue::Var(_) | ExprValue::ArrayAccess { .. }) {
            self.add_error(SemaError::InvalidAssignmentTarget, lvalue.span);
            return None;
        }
        let lvalue = self.analyze_expr(lvalue)?;
        if lvalue.ty.is_array_like() {
            self.add_error(SemaError::InvalidAssignmentTarget, lvalue.span);
            return None;
        }
        let rvalue_span = rvalue.span;
        let rvalue = self.analyze_expr(rvalue)?;
        if rvalue.ty == SemaType::Void {
            self.add_error(SemaError::InvalidOperand(SemaType::Void), rvalue_span);
            return None;
        }
        if !self.type_matches(&lvalue.ty, &rvalue.ty) {
            self.add_error(SemaError::TypeMismatch(lvalue.ty.clone(), rvalue.ty.clone()), span);
            return None;
        }
        Some(HirExpr {
            ty: lvalue.ty.clone(),
            value: HirExprValue::Assign {
                lvalue: Box::new(lvalue),
                rvalue: Box::new(rvalue),
            },
            span,
        })
    }

    fn analyze_binary_expr(&mut self, lhs: Expr, op: BinaryOp, rhs: Expr, span: Span) -> Option<HirExpr> {
        let lhs_span = lhs.span;
        let rhs_span = rhs.span;
        let lhs = self.analyze_expr(lhs)?;
        let rhs = self.analyze_expr(rhs)?;
        if !lhs.ty.is_scalar() {
            self.add_error(SemaError::InvalidOperand(lhs.ty.clone()), lhs_span);
            return None;
        }
        if !rhs.ty.is_scalar() {
            self.add_error(SemaError::InvalidOperand(rhs.ty.clone()), rhs_span);
            return None;
        }

        let result_ty = if op.is_logical() {
            SemaType::I1
        } else {
            match SemaType::get_elevate_result(&lhs.ty, &rhs.ty) {
                Some(_) if op.is_cmp() => SemaType::I1,
                Some(ty) => ty,
                None => {
                    self.add_error(SemaError::IncompatiableOperand(lhs.ty.clone(), rhs.ty.clone()), lhs_span);
                    self.add_error(SemaError::IncompatiableOperand(lhs.ty.clone(), rhs.ty.clone()), rhs_span);
                    return None;
                }
            }
        };

        Some(HirExpr {
            value: HirExprValue::BinaryOp {
                lhs: Box::new(lhs),
                op,
                rhs: Box::new(rhs),
            },
            ty: result_ty,
            span,
        })
    }

    fn analyze_func_call_expr(&mut self, func_name: String, args: Vec<Expr>, span: Span) -> Option<HirExpr> {
        let func_id = match self.function_map.get(&func_name).cloned() {
            Some(func_id) => func_id,
            None => {
                self.add_error(SemaError::FunctionNotFound(func_name), span);
                return None;
            }
        };
        let func_def = self.functions[func_id.0].clone();

        if args.len() < func_def.parameter_types.len() {
            self.add_error(SemaError::TooFewArguments(func_def.parameter_types.len(), args.len()), span);
            return None;
        }
        if args.len() > func_def.parameter_types.len() {
            self.add_error(SemaError::TooManyArguments(func_def.parameter_types.len(), args.len()), span);
            return None;
        }

        let mut has_error = false;
        for parameter_type in &func_def.parameter_types {
            if matches!(parameter_type, SemaType::Void) {
                self.add_error(SemaError::InvalidParameterType(SemaType::Void), span);
                has_error = true;
            }
        }
        if has_error {
            return None;
        }

        let mut hir_args = vec![];
        for (i, arg) in args.into_iter().enumerate() {
            let arg = self.analyze_expr(arg)?;
            let parameter_type = func_def.parameter_types[i].clone();
            if !self.type_matches(&parameter_type, &arg.ty) {
                self.add_error(SemaError::TypeMismatch(parameter_type, arg.ty.clone()), span);
                has_error = true;
                continue;
            }
            hir_args.push(arg);
        }
        if has_error {
            return None;
        }

        Some(HirExpr {
            value: HirExprValue::FuncCall(func_id, hir_args),
            ty: func_def.return_type,
            span,
        })
    }

    fn analyze_array_access_expr(&mut self, array: Expr, index: Expr, span: Span) -> Option<HirExpr> {
        let index_span = index.span;
        let array = self.analyze_expr(array)?;
        let index = self.analyze_expr(index)?;
        if !index.ty.is_scalar() {
            self.add_error(SemaError::InvalidOperand(index.ty.clone()), index_span);
            return None;
        }
        let ty = match array.ty.indexed_type() {
            Some(ty) => ty,
            None => {
                self.add_error(SemaError::NotSubscriptable, span);
                return None;
            }
        };
        Some(HirExpr {
            value: HirExprValue::ArrayAccess {
                array: Box::new(array),
                index: Box::new(index),
            },
            ty,
            span,
        })
    }

    fn type_matches(&self, expected: &SemaType, actual: &SemaType) -> bool {
        if expected == actual {
            return true;
        }
        matches!((expected, actual), (SemaType::Ptr(expected_elem), SemaType::Array(_, _)) if actual.indexed_type().is_some_and(|ty| &ty == expected_elem.as_ref()))
    }
}