//! AST to MIR lowering

use super::builder::MirBuilder;
use super::types::{
    BinOp, BlockId, Constant, Mutability, Operand, Place, Program, Rvalue, Type, UnOp,
};
use crate::frontend::ast::{
    BinaryOp as AstBinOp, Expr, ExprKind, Literal, Param, Type as AstType, UnaryOp as AstUnOp,
};
use anyhow::{anyhow, Result};
use std::collections::HashMap;

/// Context for lowering AST to MIR
pub struct LoweringContext {
    /// MIR builder
    builder: MirBuilder,
    /// Type environment for expressions
    #[allow(dead_code)]
    type_env: HashMap<String, Type>,
    /// Current block being built
    current_block: Option<BlockId>,
}

impl LoweringContext {
    /// Create a new lowering context
    #[must_use]
    pub fn new() -> Self {
        Self {
            builder: MirBuilder::new(),
            type_env: HashMap::new(),
            current_block: None,
        }
    }

    /// Lower an expression to MIR
    ///
    /// # Errors
    ///
    /// Returns an error if the expression cannot be lowered to MIR
    /// # Errors
    ///
    /// Returns an error if the operation fails
    pub fn lower_expr(&mut self, expr: &Expr) -> Result<Program> {
        match &expr.kind {
            ExprKind::Function {
                name,
                params,
                return_type,
                body,
                ..
            } => self.lower_function(name, params, return_type.as_ref(), body),
            _ => {
                // For non-function expressions, create a main function
                self.lower_main_expr(expr)
            }
        }
    }

    /// Lower a function expression
    fn lower_function(
        &mut self,
        name: &str,
        params: &[Param],
        return_type: Option<&AstType>,
        body: &Expr,
    ) -> Result<Program> {
        let func_name = name.to_string();
        let ret_ty = return_type.map_or(Type::Unit, |t| self.ast_to_mir_type(t));

        self.builder.start_function(func_name.clone(), ret_ty);

        // Add parameters
        for param in params {
            let ty = self.ast_to_mir_type(&param.ty);
            self.builder.add_param(param.name(), ty);
        }

        // Create entry block
        let entry = self.builder.new_block();
        self.current_block = Some(entry);

        // Lower function body
        let result = self.lower_expr_to_operand(body)?;

        // Return the result
        self.builder.return_(entry, Some(result));

        let function = self
            .builder
            .finish_function()
            .ok_or_else(|| anyhow!("Failed to finish function"))?;

        let mut functions = HashMap::new();
        functions.insert(func_name.clone(), function);

        Ok(Program {
            functions,
            entry: func_name,
        })
    }

    /// Lower a main expression (wrap in main function)
    fn lower_main_expr(&mut self, expr: &Expr) -> Result<Program> {
        self.builder.start_function("main".to_string(), Type::Unit);

        let entry = self.builder.new_block();
        self.current_block = Some(entry);

        // Lower the expression
        let _result = self.lower_expr_to_operand(expr)?;

        // Main function returns unit
        self.builder
            .return_(entry, Some(Operand::Constant(Constant::Unit)));

        let function = self
            .builder
            .finish_function()
            .ok_or_else(|| anyhow!("Failed to finish main function"))?;

        let mut functions = HashMap::new();
        functions.insert("main".to_string(), function);

        Ok(Program {
            functions,
            entry: "main".to_string(),
        })
    }

    /// Lower an expression to an operand
    fn lower_expr_to_operand(&mut self, expr: &Expr) -> Result<Operand> {
        let block = self
            .current_block
            .ok_or_else(|| anyhow!("No current block"))?;

        match &expr.kind {
            ExprKind::Literal(lit) => Ok(Operand::Constant(Self::lower_literal(lit))),
            ExprKind::Identifier(name) => {
                if let Some(local) = self.builder.get_local(name) {
                    Ok(Operand::Copy(Place::Local(local)))
                } else {
                    Err(anyhow!("Unbound variable: {}", name))
                }
            }
            ExprKind::Binary { op, left, right } => {
                let left_op = self.lower_expr_to_operand(left)?;
                let right_op = self.lower_expr_to_operand(right)?;
                let mir_op = Self::lower_binary_op(*op);

                // Create a variable for the result
                let result_ty = Self::infer_binary_result_type(*op);
                let temp = self.builder.alloc_local(result_ty, false, None);

                self.builder
                    .binary_op(block, temp, mir_op, left_op, right_op);
                Ok(Operand::Move(Place::Local(temp)))
            }
            ExprKind::Unary { op, operand } => {
                let operand_mir = self.lower_expr_to_operand(operand)?;
                let mir_op = Self::lower_unary_op(*op);

                let result_ty = Self::infer_unary_result_type(*op);
                let temp = self.builder.alloc_local(result_ty, false, None);

                self.builder.unary_op(block, temp, mir_op, operand_mir);
                Ok(Operand::Move(Place::Local(temp)))
            }
            ExprKind::Let {
                name, value, body, ..
            } => {
                // Lower the value
                let value_op = self.lower_expr_to_operand(value)?;

                // Create a local for the binding
                let local_ty = Type::I32; // Default type inference
                let local = self
                    .builder
                    .alloc_local(local_ty, false, Some(name.clone()));

                // Assign the value
                self.builder
                    .assign(block, Place::Local(local), Rvalue::Use(value_op));

                // Lower the body with the binding in scope
                self.lower_expr_to_operand(body)
            }
            ExprKind::If {
                condition,
                then_branch,
                else_branch,
            } => {
                let cond_op = self.lower_expr_to_operand(condition)?;

                let then_block = self.builder.new_block();
                let else_block = self.builder.new_block();
                let merge_block = self.builder.new_block();

                // Branch based on condition
                self.builder.branch(block, cond_op, then_block, else_block);

                // Lower then branch
                self.current_block = Some(then_block);
                let then_result = self.lower_expr_to_operand(then_branch)?;
                self.builder.goto(then_block, merge_block);

                // Lower else branch
                self.current_block = Some(else_block);
                let _else_result = if let Some(else_expr) = else_branch {
                    self.lower_expr_to_operand(else_expr)?
                } else {
                    Operand::Constant(Constant::Unit)
                };
                self.builder.goto(else_block, merge_block);

                // Create a variable for the result
                let result_ty = Type::I32; // Type inference would determine this
                let result_temp = self.builder.alloc_local(result_ty, false, None);

                // In merge block, we'd need phi nodes, but for simplicity we'll just use the then result
                self.current_block = Some(merge_block);
                self.builder.assign(
                    merge_block,
                    Place::Local(result_temp),
                    Rvalue::Use(then_result),
                );

                Ok(Operand::Move(Place::Local(result_temp)))
            }
            ExprKind::Call { func, args } => {
                let func_op = self.lower_expr_to_operand(func)?;
                let mut arg_ops = Vec::new();

                for arg in args {
                    arg_ops.push(self.lower_expr_to_operand(arg)?);
                }

                // Create a variable for the result
                let result_ty = Type::I32; // Type inference would determine this
                let result_temp = self.builder.alloc_local(result_ty, false, None);

                // Create call terminator
                let next_block = self.builder.call(block, result_temp, func_op, arg_ops);
                self.current_block = Some(next_block);

                Ok(Operand::Move(Place::Local(result_temp)))
            }
            ExprKind::Block(exprs) => {
                if exprs.is_empty() {
                    Ok(Operand::Constant(Constant::Unit))
                } else {
                    // Lower all expressions, return the last one
                    let mut result = Operand::Constant(Constant::Unit);
                    for expr in exprs {
                        result = self.lower_expr_to_operand(expr)?;
                    }
                    Ok(result)
                }
            }
            _ => {
                // For unsupported expressions, return unit for now
                Ok(Operand::Constant(Constant::Unit))
            }
        }
    }

    /// Lower a literal to a constant
    fn lower_literal(lit: &Literal) -> Constant {
        match lit {
            Literal::Integer(i) => Constant::Int(i128::from(*i), Type::I32),
            Literal::Float(f) => Constant::Float(*f, Type::F64),
            Literal::String(s) => Constant::String(s.clone()),
            Literal::Bool(b) => Constant::Bool(*b),
            Literal::Char(c) => Constant::Char(*c),
            Literal::Unit => Constant::Unit,
        }
    }

    /// Lower binary operator
    fn lower_binary_op(op: AstBinOp) -> BinOp {
        match op {
            AstBinOp::Add => BinOp::Add,
            AstBinOp::Subtract => BinOp::Sub,
            AstBinOp::Multiply => BinOp::Mul,
            AstBinOp::Divide => BinOp::Div,
            AstBinOp::Modulo => BinOp::Rem,
            AstBinOp::Power => BinOp::Pow,
            AstBinOp::Equal => BinOp::Eq,
            AstBinOp::NotEqual => BinOp::Ne,
            AstBinOp::Less => BinOp::Lt,
            AstBinOp::LessEqual => BinOp::Le,
            AstBinOp::Greater => BinOp::Gt,
            AstBinOp::GreaterEqual => BinOp::Ge,
            AstBinOp::And => BinOp::And,
            AstBinOp::Or => BinOp::Or,
            AstBinOp::NullCoalesce => BinOp::NullCoalesce,
            AstBinOp::BitwiseAnd => BinOp::BitAnd,
            AstBinOp::BitwiseOr => BinOp::BitOr,
            AstBinOp::BitwiseXor => BinOp::BitXor,
            AstBinOp::LeftShift => BinOp::Shl,
        }
    }

    /// Lower unary operator
    fn lower_unary_op(op: AstUnOp) -> UnOp {
        match op {
            AstUnOp::Negate => UnOp::Neg,
            AstUnOp::Not => UnOp::Not,
            AstUnOp::BitwiseNot => UnOp::BitNot,
            AstUnOp::Reference => UnOp::Ref,
        }
    }

    /// Convert AST Type to MIR Type
    #[allow(clippy::only_used_in_recursion)]
    fn ast_to_mir_type(&self, ast_ty: &AstType) -> Type {
        use crate::frontend::ast::TypeKind;
        match &ast_ty.kind {
            TypeKind::Named(name) => match name.as_str() {
                "bool" => Type::Bool,
                "i8" => Type::I8,
                "i16" => Type::I16,
                "i32" => Type::I32,
                "i64" => Type::I64,
                "i128" => Type::I128,
                "u8" => Type::U8,
                "u16" => Type::U16,
                "u32" => Type::U32,
                "u64" => Type::U64,
                "u128" => Type::U128,
                "f32" => Type::F32,
                "f64" => Type::F64,
                "String" => Type::String,
                "()" => Type::Unit,
                _ => Type::UserType(name.clone()),
            },
            TypeKind::Generic { base, params } => {
                match base.as_str() {
                    "Vec" if params.len() == 1 => {
                        Type::Vec(Box::new(self.ast_to_mir_type(&params[0])))
                    }
                    "Array" if params.len() == 1 => {
                        // For simplicity, treat arrays as vectors for now
                        Type::Vec(Box::new(self.ast_to_mir_type(&params[0])))
                    }
                    _ => Type::UserType(base.clone()),
                }
            }
            TypeKind::Optional(inner) => {
                // For simplicity, treat optionals as user types for now
                Type::UserType(format!("Option<{:?}>", inner.kind))
            }
            TypeKind::Function { params, ret } => {
                let param_types = params.iter().map(|p| self.ast_to_mir_type(p)).collect();
                Type::FnPtr(param_types, Box::new(self.ast_to_mir_type(ret)))
            }
            TypeKind::List(inner) => Type::Vec(Box::new(self.ast_to_mir_type(inner))),
            TypeKind::DataFrame { .. } => {
                // Map DataFrames to a user type for now
                Type::UserType("DataFrame".to_string())
            }
            TypeKind::Series { .. } => {
                // Map Series to a user type for now
                Type::UserType("Series".to_string())
            }
            TypeKind::Tuple(types) => {
                let mir_types: Vec<_> = types.iter().map(|t| self.ast_to_mir_type(t)).collect();
                Type::Tuple(mir_types)
            }
        }
    }

    /// Infer result type for binary operations
    fn infer_binary_result_type(op: AstBinOp) -> Type {
        match op {
            AstBinOp::Add
            | AstBinOp::Subtract
            | AstBinOp::Multiply
            | AstBinOp::Divide
            | AstBinOp::Modulo
            | AstBinOp::Power
            | AstBinOp::BitwiseAnd
            | AstBinOp::BitwiseOr
            | AstBinOp::BitwiseXor
            | AstBinOp::LeftShift => Type::I32,
            AstBinOp::Equal
            | AstBinOp::NotEqual
            | AstBinOp::Less
            | AstBinOp::LessEqual
            | AstBinOp::Greater
            | AstBinOp::GreaterEqual
            | AstBinOp::And
            | AstBinOp::Or => Type::Bool,
            AstBinOp::NullCoalesce => Type::I32, // For now, assume Int (could be improved)
        }
    }

    /// Infer result type for unary operations
    fn infer_unary_result_type(op: AstUnOp) -> Type {
        match op {
            AstUnOp::Negate | AstUnOp::BitwiseNot => Type::I32,
            AstUnOp::Not => Type::Bool,
            AstUnOp::Reference => Type::Ref(Box::new(Type::I32), Mutability::Immutable), // Reference creates an immutable reference
        }
    }
}

impl Default for LoweringContext {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::panic)]
mod tests {
    use super::*;
    use crate::frontend::Parser;

    #[test]
    fn test_lower_literal() -> Result<()> {
        let mut parser = Parser::new("42");
        let ast = parser.parse()?;

        let mut ctx = LoweringContext::new();
        let program = ctx.lower_expr(&ast)?;

        assert_eq!(program.entry, "main");
        assert!(program.functions.contains_key("main"));

        Ok(())
    }

    #[test]
    fn test_lower_binary_expr() -> Result<()> {
        let mut parser = Parser::new("1 + 2");
        let ast = parser.parse()?;

        let mut ctx = LoweringContext::new();
        let program = ctx.lower_expr(&ast)?;

        let main_func = &program.functions["main"];
        assert!(!main_func.blocks.is_empty());

        Ok(())
    }

    #[test]
    fn test_lower_function() -> Result<()> {
        let mut parser = Parser::new("fun add(x: i32, y: i32) -> i32 { x + y }");
        let ast = parser.parse()?;

        let mut ctx = LoweringContext::new();
        let program = ctx.lower_expr(&ast)?;

        assert!(program.functions.contains_key("add"));
        let func = &program.functions["add"];
        assert_eq!(func.params.len(), 2);

        Ok(())
    }

    #[test]
    fn test_lower_if_expr() -> Result<()> {
        let mut parser = Parser::new("if true { 1 } else { 2 }");
        let ast = parser.parse()?;

        let mut ctx = LoweringContext::new();
        let program = ctx.lower_expr(&ast)?;

        let main_func = &program.functions["main"];
        // Should have multiple blocks for if/else
        assert!(main_func.blocks.len() > 1);

        Ok(())
    }
}
