use std::fmt::Display;

use crate::ast::types::Type as AstType;
use crate::ast::types::BinaryOp as AstBinaryOp;
use crate::ir::err::IRError;

#[derive(Clone)]
pub enum VariableOrIntLit {
    Var(Variable),
    IntLit(i32),
}

impl Display for VariableOrIntLit {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            VariableOrIntLit::Var(v) => write!(f, "{}", v),
            VariableOrIntLit::IntLit(i) => write!(f, "{}", i),
        }
    }
}
#[derive(Clone)]
pub enum IRInstr {
    Declare(Variable),
    DeclareGlobal(Variable, Option<i32>),
    DefineFunc(Function, Vec<Variable>, Vec<IRInstr>),
    Entry,
    Binary(Variable, Variable, BinaryOp, Variable),
    Cmp(Variable, VariableOrIntLit, CmpOp, VariableOrIntLit),
    Unary(Variable, UnaryOp, Variable),
    Exit(Option<Variable>),
    FuncCall(Function, Vec<Variable>, Option<Variable>),
    Goto(usize),
    CondGoto(Variable, usize, usize), // condition, true label, false label
    Label(usize),
    Move(Variable, MoveRValue),
    Load(Variable, Variable),
    Store(Variable, Variable),
}

impl Display for IRInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            IRInstr::Entry => write!(f, "entry"),
            IRInstr::Binary(dest, left, op, right) => write!(f, "{} = {} {},{}", dest, op, left, right),
            IRInstr::Cmp(dest, left, op, right) => write!(f, "{} = icmp {} {},{}", dest, op, left, right),
            IRInstr::Exit(v) => if let Some(v) = v { write!(f, "exit {}", v) } else { write!(f, "exit") },
            IRInstr::FuncCall(func, args, dest) => {
                if let Some(dest) = dest {
                    write!(f, "{} = call {}", dest, func.to_call_string(args))
                } else {
                    write!(f, "call {}", func.to_call_string(args))
                }
            }
            IRInstr::Move(dest, src) => write!(f, "{} = {}", dest, src),
            IRInstr::Load(dest, addr) => write!(f, "{} = *{}", dest, addr),
            IRInstr::Store(addr, value) => write!(f, "*{} = {}", addr, value),
            IRInstr::Declare(var) => write!(f, "declare {}", var.to_decl_string()),
            IRInstr::DeclareGlobal(var, Some(init)) => write!(f, "declare {} = {}", var.to_decl_string(), init),
            IRInstr::DeclareGlobal(var, None) => write!(f, "declare {}", var.to_decl_string()),
            IRInstr::DefineFunc(func, args, body) => {
                let body_str = body.iter().map(|instr| format!("  {}", instr)).collect::<Vec<_>>().join("\n");
                write!(f, "define {} {{\n{}\n}}", func.to_decl_string(args), body_str)
            },
            IRInstr::Goto(label) => write!(f, "br label .L{}", label),
            IRInstr::CondGoto(cond, true_label, false_label) => write!(f, "bc {}, .L{}, .L{}", cond, true_label, false_label),
            IRInstr::Label(label) => write!(f, ".L{}:", label),
            IRInstr::Unary(dest, op, src) => write!(f, "{} = {} {}", dest, op, src),
        }
    }
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum IRType {
    I32,
    I1,
    Void,
    Ptr(Box<IRType>),
    Array(Box<IRType>, Vec<usize>),
}
impl IRType {
    pub fn size_in_bytes(&self) -> usize {
        match self {
            IRType::I32 => 4,
            IRType::I1 => 1,
            IRType::Void => 0,
            IRType::Ptr(_) => 4,
            IRType::Array(elem, dims) => elem.size_in_bytes() * dims.iter().product::<usize>(),
        }
    }

    pub fn is_array(&self) -> bool {
        matches!(self, IRType::Array(_, _))
    }

    pub fn is_array_param(&self) -> bool {
        matches!(self, IRType::Array(_, dims) if dims.first().is_some_and(|dim| *dim == 0))
    }

    pub fn decay_to_ptr(&self) -> Option<IRType> {
        match self {
            IRType::Array(elem, dims) => {
                let pointee = if dims.len() == 1 {
                    (**elem).clone()
                } else {
                    IRType::Array(elem.clone(), dims[1..].to_vec())
                };
                Some(IRType::Ptr(Box::new(pointee)))
            }
            _ => None,
        }
    }

    pub fn element_type(&self) -> Option<IRType> {
        match self {
            IRType::Array(elem, dims) => {
                if dims.len() == 1 {
                    Some((**elem).clone())
                } else {
                    Some(IRType::Array(elem.clone(), dims[1..].to_vec()))
                }
            }
            IRType::Ptr(elem) => Some((**elem).clone()),
            _ => None,
        }
    }

    pub fn get_elevate_result(lhs: &IRType, rhs: &IRType) -> Option<IRType> {
        if lhs == rhs {
            Some(lhs.clone())
        } else if (*lhs == IRType::I32 && *rhs == IRType::I1) || (*lhs == IRType::I1 && *rhs == IRType::I32) {
            Some(IRType::I32)
        } else {
            None
        }
    }
}
impl Display for IRType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            IRType::I32 => write!(f, "i32"),
            IRType::Void => write!(f, "void"),
            IRType::I1 => write!(f, "i1"),
            IRType::Ptr(elem) => write!(f, "{}*", elem),
            IRType::Array(elem, dims) => {
                write!(f, "{}", elem)?;
                for dim in dims {
                    write!(f, "[{}]", dim)?;
                }
                Ok(())
            }
        }
    }
}
impl IRType {
    pub fn base_type_and_dims(&self) -> (&IRType, &[usize]) {
        match self {
            IRType::Array(elem, dims) => (elem.as_ref(), dims.as_slice()),
            _ => (self, &[]),
        }
    }

    pub fn scalar_base_type(&self) -> &IRType {
        match self {
            IRType::Array(elem, _) => elem.scalar_base_type(),
            IRType::Ptr(elem) => elem.scalar_base_type(),
            _ => self,
        }
    }
}
impl From<AstType> for IRType {
    fn from(ast_type: AstType) -> Self {
        match ast_type {
            AstType::Int => IRType::I32,
            AstType::Void => IRType::Void,
        }
    }
}
#[derive(Clone)]
pub enum MoveRValue {
    Var(Variable),
    ConstInt(i32),
}
impl Display for MoveRValue {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MoveRValue::Var(v) => write!(f, "{}", v),
            MoveRValue::ConstInt(i) => write!(f, "{}", i),
        }
    }
}
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]
pub enum VariableType {
    Global,
    ParamTemp(usize),
    Local,
    Temp,
}
#[derive(Clone)]
pub struct Variable {
    // pub name: String,
    pub index: usize,
    pub var_type: VariableType,
    pub data_type: IRType,
}
impl PartialEq for Variable {
    fn eq(&self, other: &Self) -> bool {
        self.index == other.index && self.var_type == other.var_type
    }
}
impl Eq for Variable {}
impl PartialOrd for Variable {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(match self.index.cmp(&other.index) {
            std::cmp::Ordering::Equal => self.var_type.cmp(&other.var_type),
            ord => ord,
        })
    }
}
impl Ord for Variable {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        match self.index.cmp(&other.index) {
            std::cmp::Ordering::Equal => self.var_type.cmp(&other.var_type),
            ord => ord,
        }
    }
}
impl Display for Variable {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let prefix = match self.var_type {
            VariableType::Global => "@g",
            VariableType::Local => "%l",
            VariableType::Temp => "%t",
            VariableType::ParamTemp(_) => "%t",
        };
        write!(f, "{}{}", prefix, self.index)
    }
}
impl Variable {
    pub fn to_decl_string(&self) -> String {
        match &self.data_type {
            IRType::Array(_, _) => {
                let (base_type, dims) = self.data_type.base_type_and_dims();
                let dims_str = dims.iter().map(|dim| format!("[{}]", dim)).collect::<Vec<_>>().join("");
                format!("{} {}{}", base_type, self, dims_str)
            }
            IRType::Ptr(_) => format!("{}* {}", self.data_type.scalar_base_type(), self),
            _ => format!("{} {}", self.data_type, self),
        }
    }

    pub fn to_typed_string_with_type(&self, ty: &IRType) -> String {
        match ty {
            IRType::Array(_, _) => {
                let (base_type, dims) = ty.base_type_and_dims();
                let dims_str = dims.iter().map(|dim| format!("[{}]", dim)).collect::<Vec<_>>().join("");
                format!("{} {}{}", base_type, self, dims_str)
            }
            IRType::Ptr(_) => format!("{}* {}", ty.scalar_base_type(), self),
            _ => format!("{} {}", ty, self),
        }
    }
}
#[derive(Debug, Clone)]
pub struct Function {
    pub name: String,
    pub return_type: IRType,
    pub parameter_types: Vec<IRType>,
    // pub parameters: Vec<Variable>,
}
impl Function {
    pub fn to_call_string(&self, args: &Vec<Variable>) -> String {
        assert!(args.len() == self.parameter_types.len());
        let args_str = args.iter().zip(self.parameter_types.iter()).map(|(arg, param)| arg.to_typed_string_with_type(param)).collect::<Vec<_>>().join(", ");
        format!("{} @{}({})", self.return_type, self.name, args_str)
    }

    pub fn to_decl_string(&self, args: &Vec<Variable>) -> String {
        let params_str = args.iter().zip(self.parameter_types.iter()).map(|(arg, param_type)| {
            let (base_type, dims) = param_type.base_type_and_dims();
            let dims_str = dims.iter().map(|dim| format!("[{}]", dim)).collect::<Vec<_>>().join("");
            format!("{} {}{}", base_type, arg, dims_str)
        }).collect::<Vec<_>>().join(", ");
        format!("{} @{}({})", self.return_type, self.name, params_str)
    }
}
#[derive(Clone)]
pub enum BinaryOp {
    Add,
    Sub,
    Mul,
    Div,
    Mod,
}
#[derive(Clone)]
pub enum CmpOp {
    Le,
    Lt,
    Gt,
    Ge,
    Ne,
    Eq,
}
#[derive(Clone)]
pub enum UnaryOp {
    Neg
}

impl Display for BinaryOp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let op_str = match self {
            BinaryOp::Add => "add",
            BinaryOp::Sub => "sub",
            BinaryOp::Mul => "mul",
            BinaryOp::Div => "div",
            BinaryOp::Mod => "mod",
            // BinaryOp::Le => "cmp le",
            // BinaryOp::Lt => "cmp lt",
            // BinaryOp::Gt => "cmp gt",
            // BinaryOp::Ge => "cmp ge",
            // BinaryOp::Ne => "cmp ne",
            // BinaryOp::Eq => "cmp eq",
        };
        write!(f, "{}", op_str)
    }
}

impl Display for CmpOp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let op_str = match self {
            CmpOp::Le => "le",
            CmpOp::Lt => "lt",
            CmpOp::Gt => "gt",
            CmpOp::Ge => "ge",
            CmpOp::Ne => "ne",
            CmpOp::Eq => "eq",
        };
        write!(f, "{}", op_str)
    }
}

impl Display for UnaryOp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let op_str = match self {
            UnaryOp::Neg => "neg",
        };
        write!(f, "{}", op_str)
    }
}
impl From<AstBinaryOp> for BinaryOp {
    fn from(ast_op: AstBinaryOp) -> Self {
        match ast_op {
            AstBinaryOp::Add => BinaryOp::Add,
            AstBinaryOp::Sub => BinaryOp::Sub,
            AstBinaryOp::Mul => BinaryOp::Mul,
            AstBinaryOp::Div => BinaryOp::Div,
            AstBinaryOp::Mod => BinaryOp::Mod,
            // AstBinaryOp::Equal => BinaryOp::Eq,
            // AstBinaryOp::NotEqual => BinaryOp::Ne,
            // AstBinaryOp::Less => BinaryOp::Lt,
            // AstBinaryOp::LessEqual => BinaryOp::Le,
            // AstBinaryOp::Greater => BinaryOp::Gt,
            // AstBinaryOp::GreaterEqual => BinaryOp::Ge,
            _ => unreachable!(),
        }
    }
}
impl From<AstBinaryOp> for CmpOp {
    fn from(ast_op: AstBinaryOp) -> Self {
        match ast_op {
            AstBinaryOp::Equal => CmpOp::Eq,
            AstBinaryOp::NotEqual => CmpOp::Ne,
            AstBinaryOp::Less => CmpOp::Lt,
            AstBinaryOp::LessEqual => CmpOp::Le,
            AstBinaryOp::Greater => CmpOp::Gt,
            AstBinaryOp::GreaterEqual => CmpOp::Ge,
            _ => unreachable!(),
        }
    }
}
