use std::collections::HashMap;

use cranelift_codegen::{
    Context,
    control::ControlPlane,
    ir::{self, AbiParam, InstBuilder, types},
    settings::{self, Configurable},
};
use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext, Variable};
use cranelift_module::{Linkage, Module, default_libcall_names};
use cranelift_object::{ObjectBuilder, ObjectModule};

use crate::frontend::{
    ast::{BinaryOp, UnaryOp},
    sema::Ty,
};
use crate::middle::mir::*;

/// The backend responsible for lowering a `TypedModule` into Cranelift IR and
/// generating native machine code object files.
pub struct CraneliftBackend {
    ctx: Context,
    builder_context: FunctionBuilderContext,
    module: ObjectModule,
}

impl CraneliftBackend {
    /// Initializes the Cranelift backend with the host's native instruction set architecture (ISA),
    /// enabling speed optimizations and position-independent code (PIC) generation.
    pub fn new() -> Self {
        let mut flag_builder = settings::builder();
        flag_builder.set("is_pic", "true").unwrap();
        flag_builder.set("opt_level", "speed").unwrap();

        let isa_builder = cranelift_native::builder().unwrap_or_else(|msg| {
            panic!("host machine is not supported: {}", msg);
        });

        let isa = isa_builder
            .finish(settings::Flags::new(flag_builder))
            .unwrap();

        let builder = ObjectBuilder::new(isa, "module", default_libcall_names()).unwrap();

        Self {
            ctx: Context::new(),
            builder_context: FunctionBuilderContext::new(),
            module: ObjectModule::new(builder),
        }
    }

    /// Compiles a MIR module into native object code.
    ///
    /// Returns a tuple containing the generated Cranelift IR (as a human-readable string) and the assembled object file bytes.
    pub fn compile_module(mut self, module: &MirModule) -> (String, Vec<u8>) {
        let mut ir_output = String::new();

        for func in &module.functions {
            self.compile_function(func);

            // Run Cranelift's optimization passes before emitting the text IR
            let mut ctrl_plane = ControlPlane::default();
            self.ctx
                .optimize(self.module.isa(), &mut ctrl_plane)
                .unwrap();

            ir_output.push_str(&format!(
                "; Function: {}\n{}",
                func.name,
                self.ctx.func.to_string()
            ));
            ir_output.push('\n');

            let func_id = self
                .module
                .declare_function(&func.name, Linkage::Export, &self.ctx.func.signature)
                .unwrap();

            self.module.define_function(func_id, &mut self.ctx).unwrap();
            self.module.clear_context(&mut self.ctx);
        }

        let obj_bytes = self.module.finish().emit().unwrap();
        (ir_output, obj_bytes)
    }

    /// Lowers a single MIR function into Cranelift IR.
    fn compile_function(&mut self, func: &MirFunction) {
        let mut sig = self.module.make_signature();

        for param_id in &func.params {
            let param_ty = &func.locals[param_id.0].ty;
            sig.params.push(AbiParam::new(Self::lower_type(param_ty)));
        }

        if func.return_type != Ty::Unit {
            sig.returns
                .push(AbiParam::new(Self::lower_type(&func.return_type)));
        }

        self.ctx.func.signature = sig;
        self.ctx.func.name = ir::UserFuncName::user(0, 0);

        let mut builder = FunctionBuilder::new(&mut self.ctx.func, &mut self.builder_context);

        let mut block_map = HashMap::new();
        for block in &func.blocks {
            block_map.insert(block.id, builder.create_block());
        }

        let mut var_map = HashMap::new();
        for local in &func.locals {
            let var = builder.declare_var(Self::lower_type(&local.ty));
            var_map.insert(local.id, var);
        }

        let mut trans = FunctionTranslator {
            builder,
            var_map,
            block_map,
            locals: &func.locals,
        };

        if let Some(first_block) = func.blocks.first() {
            let entry_block = trans.block_map[&first_block.id];
            trans
                .builder
                .append_block_params_for_function_params(entry_block);
        }

        for (i, block) in func.blocks.iter().enumerate() {
            let cl_block = trans.block_map[&block.id];
            trans.builder.switch_to_block(cl_block);

            // Retrieve function arguments explicitly if this is the entry block
            if i == 0 {
                for (j, param_id) in func.params.iter().enumerate() {
                    let val = trans.builder.block_params(cl_block)[j];
                    trans.builder.def_var(trans.var_map[param_id], val);
                }
            }

            for stmt in &block.statements {
                trans.translate_stmt(stmt);
            }

            trans.translate_terminator(&block.terminator);
        }

        // Seal all blocks now that all branches are translated and predecessors are known
        for cl_block in trans.block_map.values() {
            trans.builder.seal_block(*cl_block);
        }

        trans.builder.finalize();
    }

    /// Maps our semantic types (`Ty`) to Cranelift's internal types (`ir::Type`).
    fn lower_type(ty: &Ty) -> ir::Type {
        match ty {
            Ty::I8 | Ty::U8 => types::I8,
            Ty::I16 | Ty::U16 => types::I16,
            Ty::I32 | Ty::U32 => types::I32,
            Ty::I64 | Ty::U64 => types::I64,
            Ty::Bool => types::I8, // Booleans are represented as 8-bit integers
            _ => unimplemented!("Unsupported type for Cranelift lowering: {:?}", ty),
        }
    }
}

/// A visitor that traverses MIR basic blocks and instructions, emitting Cranelift IR instructions
/// into the current function builder.
struct FunctionTranslator<'a> {
    builder: FunctionBuilder<'a>,
    var_map: HashMap<LocalId, Variable>,
    block_map: HashMap<BlockId, ir::Block>,
    locals: &'a [LocalDecl],
}

impl<'a> FunctionTranslator<'a> {
    fn translate_stmt(&mut self, stmt: &Statement) {
        match &stmt.kind {
            StatementKind::Assign(local_id, rvalue) => {
                let val = self.translate_rvalue(rvalue);
                let var = self.var_map[local_id];
                self.builder.def_var(var, val);
            }
        }
    }

    fn translate_terminator(&mut self, term: &Terminator) {
        match &term.kind {
            TerminatorKind::Goto { target } => {
                self.builder.ins().jump(self.block_map[target], &[]);
            }
            TerminatorKind::CondBranch {
                cond,
                target_true,
                target_false,
            } => {
                let cond_val = self.translate_operand(cond);
                self.builder.ins().brif(
                    cond_val,
                    self.block_map[target_true],
                    &[],
                    self.block_map[target_false],
                    &[],
                );
            }
            TerminatorKind::Return { value } => {
                if let Some(op) = value {
                    let val = self.translate_operand(op);
                    self.builder.ins().return_(&[val]);
                } else {
                    self.builder.ins().return_(&[]);
                }
            }
            TerminatorKind::Unreachable => {
                self.builder.ins().trap(ir::TrapCode::user(5).unwrap());
            }
        }
    }

    fn get_operand_type(&self, op: &Operand) -> Ty {
        match op {
            Operand::Copy(local_id) => self.locals[local_id.0].ty.clone(),
            Operand::Constant(ConstantValue::Integer(_, ty)) => ty.clone(),
            Operand::Constant(ConstantValue::Boolean(_)) => Ty::Bool,
        }
    }

    fn translate_operand(&mut self, op: &Operand) -> ir::Value {
        match op {
            Operand::Copy(local_id) => {
                let var = self.var_map[local_id];
                self.builder.use_var(var)
            }
            Operand::Constant(ConstantValue::Integer(val, ty)) => {
                let cl_ty = CraneliftBackend::lower_type(ty);
                self.builder.ins().iconst(cl_ty, *val as i64)
            }
            Operand::Constant(ConstantValue::Boolean(val)) => self
                .builder
                .ins()
                .iconst(types::I8, if *val { 1 } else { 0 }),
        }
    }

    fn translate_rvalue(&mut self, rvalue: &Rvalue) -> ir::Value {
        match rvalue {
            Rvalue::Use(op) => self.translate_operand(op),
            Rvalue::UnaryOp(op, inner) => {
                let inner_val = self.translate_operand(inner);
                match op {
                    UnaryOp::Neg => self.builder.ins().ineg(inner_val),
                    UnaryOp::Not => {
                        let ty = self.get_operand_type(inner);
                        let cl_ty = CraneliftBackend::lower_type(&ty);
                        let zero = self.builder.ins().iconst(cl_ty, 0);
                        self.builder
                            .ins()
                            .icmp(ir::condcodes::IntCC::Equal, inner_val, zero)
                    }
                }
            }
            Rvalue::BinaryOp(op, lhs, rhs) => {
                let lhs_val = self.translate_operand(lhs);
                let rhs_val = self.translate_operand(rhs);

                let ty = self.get_operand_type(lhs);
                let is_signed = matches!(ty, Ty::I8 | Ty::I16 | Ty::I32 | Ty::I64);

                match op {
                    BinaryOp::Add => self.builder.ins().iadd(lhs_val, rhs_val),
                    BinaryOp::Sub => self.builder.ins().isub(lhs_val, rhs_val),
                    BinaryOp::Mul => self.builder.ins().imul(lhs_val, rhs_val),
                    BinaryOp::Div => {
                        if is_signed {
                            self.builder.ins().sdiv(lhs_val, rhs_val)
                        } else {
                            self.builder.ins().udiv(lhs_val, rhs_val)
                        }
                    }
                    BinaryOp::Rem => {
                        if is_signed {
                            self.builder.ins().srem(lhs_val, rhs_val)
                        } else {
                            self.builder.ins().urem(lhs_val, rhs_val)
                        }
                    }
                    BinaryOp::Eq => {
                        self.builder
                            .ins()
                            .icmp(ir::condcodes::IntCC::Equal, lhs_val, rhs_val)
                    }
                    BinaryOp::Neq => {
                        self.builder
                            .ins()
                            .icmp(ir::condcodes::IntCC::NotEqual, lhs_val, rhs_val)
                    }
                    BinaryOp::Lt => {
                        let cc = if is_signed {
                            ir::condcodes::IntCC::SignedLessThan
                        } else {
                            ir::condcodes::IntCC::UnsignedLessThan
                        };
                        self.builder.ins().icmp(cc, lhs_val, rhs_val)
                    }
                    BinaryOp::Le => {
                        let cc = if is_signed {
                            ir::condcodes::IntCC::SignedLessThanOrEqual
                        } else {
                            ir::condcodes::IntCC::UnsignedLessThanOrEqual
                        };
                        self.builder.ins().icmp(cc, lhs_val, rhs_val)
                    }
                    BinaryOp::Gt => {
                        let cc = if is_signed {
                            ir::condcodes::IntCC::SignedGreaterThan
                        } else {
                            ir::condcodes::IntCC::UnsignedGreaterThan
                        };
                        self.builder.ins().icmp(cc, lhs_val, rhs_val)
                    }
                    BinaryOp::Ge => {
                        let cc = if is_signed {
                            ir::condcodes::IntCC::SignedGreaterThanOrEqual
                        } else {
                            ir::condcodes::IntCC::UnsignedGreaterThanOrEqual
                        };
                        self.builder.ins().icmp(cc, lhs_val, rhs_val)
                    }
                }
            }
        }
    }
}