#![allow(dead_code)]

use std::collections::HashMap;

use crate::{
    ast::{FloatingType, IntegralType, Node as AstNode, Tag, Type},
    parser::Tree,
    writeln_indented,
};

type Node = u32;

enum Inst {
    Label(String),
    Constant(Value),
    Parameter,
    Add { lhs: Node, rhs: Node },
    Sub { lhs: Node, rhs: Node },
    Div { lhs: Node, rhs: Node },
    Mul { lhs: Node, rhs: Node },
    Rem { lhs: Node, rhs: Node },
    BitAnd { lhs: Node, rhs: Node },
    BitOr { lhs: Node, rhs: Node },
    BitXOr { lhs: Node, rhs: Node },
    Negate { lhs: Node },
    ReturnValue { lhs: Node },
    Return,
    Alloc { size: u32, align: u32 },
    AddressOf(Node),
    Load { source: Node },
    Store { dest: Node, source: Node },
}

enum Value {
    Int { kind: IntegralType, bits: u64 },
    Float { kind: FloatingType, bits: u64 },
}

impl core::fmt::Display for Value {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Value::Int { kind, bits } => {
                write!(f, "{kind} {}", bits)
            }
            Value::Float { kind, bits } => {
                write!(f, "{kind} {{{}}}", bits)
            }
        }
    }
}

struct IRBuilder<'tree, 'ir> {
    ir: &'ir mut IR,
    tree: &'tree mut Tree,
    type_map: HashMap<AstNode, Type>,
    lookup: HashMap<AstNode, Node>,
}

impl<'tree, 'ir> IRBuilder<'tree, 'ir> {
    fn new(ir: &'ir mut IR, tree: &'tree mut Tree) -> Self {
        Self {
            ir,
            tree,
            type_map: HashMap::new(),
            lookup: HashMap::new(),
        }
    }

    fn visit(&mut self, node: AstNode) -> Node {
        match &self.tree.nodes[node].clone() {
            Tag::FunctionDecl { proto, body } => {
                self.visit(*proto);
                self.tree.st.into_child(node);
                let value = self.visit(*body);
                // TODO: return value of body expression
                let node = if value != !0 {
                    let return_type = {
                        match self.tree.nodes.get_node(*proto) {
                            Tag::FunctionProto { return_type, .. } => *return_type,
                            _ => unreachable!(),
                        }
                    };
                    self.type_check(return_type, *body);
                    self.ir.push(Inst::ReturnValue { lhs: value })
                } else {
                    !0
                };

                self.tree.st.into_parent();
                node
            }
            Tag::FunctionProto {
                parameters, name, ..
            } => {
                parameters.map(|p| self.visit(p));

                let label = self.ir.push(Inst::Label(
                    self.tree.nodes.get_ident_str(*name).unwrap().to_string(),
                ));

                self.lookup.insert(node, label);

                label
            }
            Tag::ParameterList { parameters } => {
                for param in parameters {
                    self.visit(*param);
                }
                !0
            }
            Tag::Parameter { .. } => {
                let param = self.ir.push(Inst::Parameter);

                self.lookup.insert(node, param);
                param
            }
            Tag::Block {
                statements,
                trailing_expr,
            } => {
                for stmt in statements {
                    self.visit(*stmt);
                }

                if let Some(expr) = trailing_expr {
                    self.visit(*expr)
                } else {
                    !0
                }
            }
            Tag::VarDecl { .. } => {
                let ty = self.tree.type_of_node(node);
                let alloca = self.ir.push(Inst::Alloc {
                    size: ty.size_of(),
                    align: ty.align_of(),
                });
                self.lookup.insert(node, alloca);
                alloca
            }
            Tag::ReturnStmt { expr } => {
                if let Some(expr) = expr {
                    let expr = self.visit(*expr);
                    self.ir.push(Inst::ReturnValue { lhs: expr })
                } else {
                    self.ir.push(Inst::Return)
                }
            }
            Tag::ExprStmt { expr } => self.visit(*expr),
            Tag::Deref { lhs } => {
                let lhs = self.visit(*lhs);
                self.ir.push(Inst::Load { source: lhs })
            }
            Tag::IntegralConstant { bits, ty } => self.ir.push(Inst::Constant(Value::Int {
                kind: *ty,
                bits: *bits,
            })),
            Tag::FloatingConstant { bits, ty } => self.ir.push(Inst::Constant(Value::Float {
                kind: *ty,
                bits: *bits,
            })),
            Tag::Assign { lhs, rhs } => {
                self.type_check(*lhs, *rhs);

                let lhs = self.visit(*lhs);
                let rhs = self.visit(*rhs);
                self.ir.push(Inst::Store {
                    dest: lhs,
                    source: rhs,
                })
            }
            Tag::Add { lhs, rhs } => {
                let ty = self.type_check(*lhs, *rhs);
                if !ty.can_add_sub() {
                    eprintln!("add is not available for type {ty:?}");
                }

                let lhs = self.visit(*lhs);
                let rhs = self.visit(*rhs);
                self.ir.push(Inst::Add { lhs, rhs })
            }
            Tag::Sub { lhs, rhs } => {
                let ty = self.type_check(*lhs, *rhs);
                if !ty.can_add_sub() {
                    eprintln!("sub is not available for type {ty:?}");
                }

                let lhs = self.visit(*lhs);
                let rhs = self.visit(*rhs);
                self.ir.push(Inst::Sub { lhs, rhs })
            }
            Tag::Mul { lhs, rhs } => {
                let ty = self.type_check(*lhs, *rhs);
                if !ty.can_mul_div_rem() {
                    eprintln!("mul is not available for type {ty:?}");
                }

                let lhs = self.visit(*lhs);
                let rhs = self.visit(*rhs);
                self.ir.push(Inst::Mul { lhs, rhs })
            }
            Tag::Div { lhs, rhs } => {
                let ty = self.type_check(*lhs, *rhs);
                if !ty.can_mul_div_rem() {
                    eprintln!("div is not available for type {ty:?}");
                }

                let lhs = self.visit(*lhs);
                let rhs = self.visit(*rhs);
                self.ir.push(Inst::Div { lhs, rhs })
            }
            Tag::Rem { lhs, rhs } => {
                let ty = self.type_check(*lhs, *rhs);
                if !ty.can_mul_div_rem() {
                    eprintln!("rem is not available for type {ty:?}");
                }

                let lhs = self.visit(*lhs);
                let rhs = self.visit(*rhs);
                self.ir.push(Inst::Rem { lhs, rhs })
            }
            Tag::BitAnd { lhs, rhs } => {
                let ty = self.type_check(*lhs, *rhs);
                if !ty.can_bitxor_and_or() {
                    eprintln!("bitand is not available for type {ty:?}");
                }

                let lhs = self.visit(*lhs);
                let rhs = self.visit(*rhs);
                self.ir.push(Inst::BitAnd { lhs, rhs })
            }
            Tag::BitOr { lhs, rhs } => {
                let ty = self.type_check(*lhs, *rhs);
                if !ty.can_bitxor_and_or() {
                    eprintln!("bitor is not available for type {ty:?}");
                }

                let lhs = self.visit(*lhs);
                let rhs = self.visit(*rhs);
                self.ir.push(Inst::BitOr { lhs, rhs })
            }
            Tag::BitXOr { lhs, rhs } => {
                let ty = self.type_check(*lhs, *rhs);
                if !ty.can_bitxor_and_or() {
                    eprintln!("bitxor is not available for type {ty:?}");
                }

                let lhs = self.visit(*lhs);
                let rhs = self.visit(*rhs);
                self.ir.push(Inst::BitXOr { lhs, rhs })
            }
            Tag::Negate { lhs } => {
                let ty = self.tree.type_of_node(*lhs);
                if !ty.can_negate() {
                    eprintln!("negation is not available for type {ty:?}");
                }

                let lhs = self.visit(*lhs);
                self.ir.push(Inst::Negate { lhs })
            }
            Tag::DeclRef(decl) => *self.lookup.get(decl).expect("declref not in lookup map"),
            Tag::Ref { lhs } => {
                let lhs = self.visit(*lhs);
                self.ir.push(Inst::AddressOf(lhs))
            }
            _ => {
                dbg!(&self.tree.nodes[node]);
                todo!()
            }
        }
    }

    fn type_check(&self, lhs: AstNode, rhs: AstNode) -> Type {
        let t_lhs = self.tree.type_of_node(lhs);
        let t_rhs = self.tree.type_of_node(rhs);
        if t_lhs != t_rhs {
            eprintln!("incompatible types {t_lhs:?} and {t_rhs:?}!");
        }
        t_lhs
    }
}

struct IR {
    nodes: Vec<Inst>,
}

impl IR {
    pub fn new() -> Self {
        Self { nodes: Vec::new() }
    }

    fn push(&mut self, inst: Inst) -> u32 {
        let node = self.nodes.len() as u32;
        self.nodes.push(inst);
        node
    }

    pub fn build(&mut self, tree: &mut Tree, ast_node: crate::ast::Node) {
        let mut builder = IRBuilder::new(self, tree);

        builder.visit(ast_node);
    }

    fn render_node<W: core::fmt::Write>(
        &self,
        w: &mut W,
        node: Node,
        indent: u32,
    ) -> core::fmt::Result {
        match &self.nodes[node as usize] {
            Inst::Label(label) => {
                writeln_indented!(indent - 1, w, "%{} = {label}:", node)?;
            }
            Inst::Parameter => {
                writeln_indented!(indent, w, "%{} = Param", node)?;
            }
            Inst::Constant(value) => {
                writeln_indented!(indent, w, "%{} = {}", node, value)?;
            }
            Inst::Add { lhs, rhs } => {
                writeln_indented!(indent, w, "%{} = %{} + %{}", node, lhs, rhs)?;
            }
            Inst::Sub { lhs, rhs } => {
                writeln_indented!(indent, w, "%{} = %{} - %{}", node, lhs, rhs)?;
            }
            Inst::Mul { lhs, rhs } => {
                writeln_indented!(indent, w, "%{} = %{} * %{}", node, lhs, rhs)?;
            }
            Inst::Div { lhs, rhs } => {
                writeln_indented!(indent, w, "%{} = %{} / %{}", node, lhs, rhs)?;
            }
            Inst::Rem { lhs, rhs } => {
                writeln_indented!(indent, w, "%{} = %{} % %{}", node, lhs, rhs)?;
            }
            Inst::BitAnd { lhs, rhs } => {
                writeln_indented!(indent, w, "%{} = %{} & %{}", node, lhs, rhs)?;
            }
            Inst::BitOr { lhs, rhs } => {
                writeln_indented!(indent, w, "%{} = %{} | %{}", node, lhs, rhs)?;
            }
            Inst::BitXOr { lhs, rhs } => {
                writeln_indented!(indent, w, "%{} = %{} ^ %{}", node, lhs, rhs)?;
            }
            Inst::Negate { lhs } => {
                writeln_indented!(indent, w, "%{} = !%{}", node, lhs)?;
            }
            Inst::ReturnValue { lhs } => {
                writeln_indented!(indent, w, "%{} = return %{}", node, lhs)?;
            }
            Inst::Return => {
                writeln_indented!(indent, w, "%{} = return", node)?;
            }
            Inst::Alloc { size, align } => {
                writeln_indented!(indent, w, "%{} = alloca {size} (algin: {align})", node)?;
            }
            Inst::AddressOf(val) => {
                writeln_indented!(indent, w, "%{} = addr %{val}", node)?;
            }
            Inst::Load { source } => {
                writeln_indented!(indent, w, "%{} = load ptr %{source}", node)?;
            }
            Inst::Store { dest, source } => {
                writeln_indented!(indent, w, "%{} = store ptr %{dest} from %{source}", node)?;
            }
        }
        Ok(())
    }

    pub fn render<W: core::fmt::Write>(&self, w: &mut W) -> core::fmt::Result {
        for node in 0..self.nodes.len() {
            self.render_node(w, node as u32, 1)?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::lexer::Tokenizer;

    use super::*;

    #[test]
    fn ir() {
        let src = "
fn main() -> u32 {
let a: u32 = 0 + 3;
let ptr_a = &a;
return *ptr_a * global;
}
let global: u32 = 42;
";
        let tokens = Tokenizer::new(src.as_bytes()).unwrap();

        let mut tree = Tree::new();
        tree.parse(tokens.iter()).unwrap();

        let mut buf = String::new();
        tree.render(&mut buf).unwrap();
        println!("{buf}");

        let mut ir = IR::new();
        let decl = *tree.global_decls.first().unwrap();
        ir.build(&mut tree, decl);
        let mut buf = String::new();
        ir.render(&mut buf).unwrap();
        println!("{buf}");
    }
}