#![allow(unused_variables)]

use std::collections::HashMap;

use crate::{
    checked::{self, Node, NodeKind},
    ir::{Block, Data, DataKind, Fn, Op},
    sym::Syms,
};

pub struct Program {
    pub fns: Vec<Fn>,
    pub data: Vec<Data>,
    pub entry: u64,
    pub sym: Syms,
}

struct Compiler {
    syms: Syms,
}

impl Compiler {
    pub fn new(syms: Syms) -> Self {
        Self { syms }
    }

    pub fn compile(self, ast: &[Node]) -> Result<Program, ()> {
        let (fns, data) = FnCompiler::new(&self.syms).compile_entry(ast)?;
        Ok(Program {
            fns,
            data,
            sym: self.syms,
            entry: 0,
        })
    }

    fn error<S: Into<String>>(&mut self, msg: S) {
        println!("ir compiler: {}", msg.into());
    }
}

struct FnCompiler<'syms> {
    syms: &'syms Syms,
    ops: Vec<Op>,
    local_count: i32,
    local_map: HashMap<u64, i32>,
    fns: Vec<Fn>,
    data: Vec<Data>,
}

impl<'syms> FnCompiler<'syms> {
    pub fn new(syms: &'syms Syms) -> Self {
        Self {
            syms,
            ops: Vec::new(),
            local_count: 0,
            local_map: HashMap::new(),
            fns: Vec::new(),
            data: Vec::new(),
        }
    }

    pub fn compile_entry(mut self, ast: &[Node]) -> Result<(Vec<Fn>, Vec<Data>), ()> {
        let mut blocks = Vec::new();
        for stmt in ast {
            blocks.append(&mut self.compile_stmt(stmt)?);
            if stmt.typ != checked::Type::Unit {
                self.ops.push(Op::Pop);
            }
        }
        self.fns.push(Fn {
            blocks,
            id: 0,
            arg_count: 0,
            local_count: self.local_count,
        });
        Ok((self.fns, self.data))
    }

    pub fn compile(mut self, node: &Node) -> Result<Fn, ()> {
        let Node {
            kind:
                NodeKind::Fn {
                    subject: _,
                    params,
                    return_typ: _,
                    body,
                    id,
                },
            ..
        } = node
        else {
            unreachable!()
        };
        let NodeKind::Block(body) = &body.kind else {
            unreachable!()
        };
        let mut blocks = Vec::new();
        for stmt in body {
            blocks.append(&mut self.compile_stmt(stmt)?);
            if stmt.typ != checked::Type::Unit {
                self.ops.push(Op::Pop);
            }
        }
        self.push_ops(&mut blocks);
        Ok(Fn {
            blocks,
            id: *id,
            arg_count: params.len() as i32,
            local_count: self.local_count,
        })
    }

    pub fn compile_stmt(&mut self, stmt: &Node) -> Result<Vec<Block>, ()> {
        let syms = self.syms.view(stmt.table_id);
        let mut blocks = Vec::new();
        match &stmt.kind {
            NodeKind::Error => return Err(()),
            NodeKind::Break => todo!(),
            NodeKind::Assign { subject, value } => todo!(),
            NodeKind::Let { subject, value } => todo!(),
            NodeKind::Fn {
                subject,
                params,
                return_typ,
                body,
                id,
            } => todo!(),
            NodeKind::Return { value } => {
                self.ops.push(Op::Return);
            }
            _ => {
                blocks.append(&mut self.compile_expr(stmt)?);
            }
        }
        if stmt.typ != checked::Type::Unit {
            self.ops.push(Op::Pop);
        }
        Ok(blocks)
    }

    fn compile_expr(&mut self, expr: &Node) -> Result<Vec<Block>, ()> {
        let syms = self.syms.view(expr.table_id);
        let mut blocks = Vec::new();
        match &expr.kind {
            NodeKind::Error => return Err(()),
            NodeKind::Id(id) => {
                let sym = syms.get(*id).unwrap();
                let local_id = self.local_map.get(&sym.uid).unwrap();
                self.ops.push(Op::LoadLocal(*local_id));
            }
            NodeKind::Int(value) => match expr.typ {
                checked::Type::I32 => self.ops.push(Op::PushI32(*value as i32)),
                checked::Type::U32 => self.ops.push(Op::PushU32(*value as u32)),
                _ => unreachable!(),
            },
            NodeKind::String(value) => {
                let id = self.data.len() as u64;
                let bytes = value.bytes().collect();
                self.data.push(Data {
                    kind: DataKind::U8(bytes),
                    id,
                });
                self.ops.push(Op::PushStaticPtr(id));
            }
            NodeKind::Group(expr) => {
                blocks.append(&mut self.compile_expr(expr)?);
            }
            NodeKind::Block(stmts) => {
                let mut last_typ = None;
                for stmt in stmts {
                    if last_typ.filter(|typ| *typ != checked::Type::Unit).is_some() {
                        self.ops.push(Op::Pop);
                    }
                    last_typ = Some(stmt.typ.clone());
                    self.compile_stmt(stmt)?;
                }
            }
            NodeKind::Call { subject, args } => {
                self.compile_expr(subject)?;
                for arg in args {
                    self.compile_expr(arg)?;
                }
                match subject.typ {
                    checked::Type::Fn {
                        id,
                        params: _,
                        return_typ: _,
                    } => {
                        self.ops.push(Op::CallFn(id, args.len() as i32));
                    }
                    _ => {
                        self.ops.push(Op::CallPtr(args.len() as i32));
                    }
                }
            }
            NodeKind::If {
                cond,
                truthy,
                falsy,
            } => {
                self.compile_expr(cond)?;
                let cond_idx = blocks.len();
                self.push_ops(&mut blocks);

                match falsy {
                    Some(falsy) => {
                        let truthy_first_idx = blocks.len() as u64;
                        let mut truthy = self.compile_expr(truthy)?;
                        let truthy_last_idx = blocks.len();
                        self.push_ops(&mut truthy);
                        blocks.append(&mut truthy);

                        let falsy_first_idx = blocks.len() as u64;
                        let mut falsy = self.compile_expr(falsy)?;
                        let falsy_last_idx = blocks.len();
                        self.push_ops(&mut falsy);
                        blocks.append(&mut falsy);

                        let after_idx = blocks.len() as u64;

                        blocks[cond_idx].ops.push(Op::JumpIfTrue(truthy_first_idx));
                        blocks[cond_idx].ops.push(Op::Jump(falsy_first_idx));
                        blocks[truthy_last_idx].ops.push(Op::Jump(after_idx));
                        blocks[falsy_last_idx].ops.push(Op::Jump(after_idx));
                    }
                    None => {
                        let truthy_first_idx = blocks.len() as u64;
                        let mut truthy = self.compile_expr(truthy)?;
                        let truthy_last_idx = blocks.len();
                        self.push_ops(&mut truthy);
                        blocks.append(&mut truthy);

                        let after_idx = blocks.len() as u64;

                        blocks[cond_idx].ops.push(Op::JumpIfTrue(truthy_first_idx));
                        blocks[cond_idx].ops.push(Op::Jump(after_idx));
                        blocks[truthy_last_idx].ops.push(Op::Jump(after_idx));
                    }
                }
            }
            NodeKind::Loop { body } => {
                let body = self.compile_expr(body)?;
                let body_idx = blocks.len();
            }
            _ => unreachable!(),
        }
        self.push_ops(&mut blocks);
        Ok(blocks)
    }

    fn push_ops(&mut self, blocks: &mut Vec<Block>) {
        let mut ops = Vec::new();
        std::mem::swap(&mut self.ops, &mut ops);
        blocks.push(Block { ops });
    }

    fn error<S: Into<String>>(&mut self, msg: S) {
        println!("ir compiler: {}", msg.into());
    }
}

#[test]
fn test_compiler() {
    use crate::checker::{Checker, IdGen};
    use crate::parser::Parser;
    use pretty_assertions::assert_eq;
    use Op::*;

    struct SeqIdGen(u64);
    impl IdGen for SeqIdGen {
        fn new() -> Self {
            Self(0)
        }

        fn gen(&mut self) -> u64 {
            let v = self.0;
            self.0 += 1;
            v
        }
    }

    let compile = |text| {
        let mut checker = Checker::<SeqIdGen>::new_with_fn_id_gen();
        let checked = checker.check(&Parser::new(text).parse());
        let syms = checker.finish();
        let compiled = Compiler::new(syms).compile(&checked);
        compiled
    };

    assert_eq!(
        compile("123;").map(|program| program.fns),
        Ok(vec![Fn {
            blocks: vec![Block {
                ops: vec![PushI32(123)]
            }],
            id: 0,
            arg_count: 0,
            local_count: 0
        }])
    );
}