#![allow(unused_variables)]

use rand::random;

use crate::{
    checked::{Node, NodeKind, Type},
    itertewls::Itertewls,
    parsed,
    sym::Syms,
    util::hash,
};

pub trait IdGen {
    fn new() -> Self;
    fn gen(&mut self) -> u64;
}

pub struct RandIdGen;
impl IdGen for RandIdGen {
    fn new() -> Self {
        Self
    }
    fn gen(&mut self) -> u64 {
        random()
    }
}

pub struct Checker<FnIdGen: IdGen = RandIdGen> {
    syms: Syms,
    fn_id_gen: FnIdGen,
}
impl Checker {
    pub fn new() -> Self {
        Self {
            syms: Syms::new(),
            fn_id_gen: RandIdGen::new(),
        }
    }
}

impl<FnIdGen: IdGen> Checker<FnIdGen> {
    pub fn new_with_fn_id_gen() -> Self {
        Self {
            syms: Syms::new(),
            fn_id_gen: FnIdGen::new(),
        }
    }

    pub fn check(&mut self, ast: &[parsed::Node]) -> Vec<Node> {
        self.fn_scan(ast);
        ast.iter().map(|stmt| self.check_expr(stmt)).collect()
    }

    fn fn_scan(&mut self, ast: &[parsed::Node]) {
        for node in ast {
            if let parsed::Node::Fn {
                subject,
                params,
                return_typ,
                body,
            } = node
            {
                let Ok(params) = self.fn_scan_params(params) else {
                    continue;
                };

                if let Some(id) = params.iter().map(|(id, _)| *id).find_first_duplicate() {
                    self.error("redefinition param");
                    continue;
                }

                let params = params.into_iter().map(|(_, param)| param).collect();

                let parsed::Node::Id(id) = subject.as_ref() else { unreachable!() };

                if self.syms.defined_locally(*id) {
                    self.error("redefinition fn");
                    continue;
                }

                let return_typ = self.check_type(return_typ);

                let typ = Type::Fn {
                    id: self.fn_id_gen.gen(),
                    params,
                    return_typ: Box::new(return_typ),
                };

                self.syms.define(*id, typ.clone());
            }
        }
    }

    fn fn_scan_params(&self, params: &[parsed::Node]) -> Result<Vec<(u64, Node)>, ()> {
        params
            .iter()
            .map(|param| {
                let parsed::Node::Param { subject, typ } = param else { unreachable!() };
                let parsed::Node::Id(id) = subject.as_ref() else { unreachable!() };
                let typ = self.check_type(typ.as_ref().ok_or(())?);
                Ok((
                    *id,
                    self.node(
                        NodeKind::Param {
                            subject: Box::new(self.node(NodeKind::Id(*id), Type::Unit)),
                            typ: Some(typ),
                        },
                        Type::Unit,
                    ),
                ))
            })
            .collect::<Result<Vec<_>, _>>()
    }

    fn check_expr(&mut self, node: &parsed::Node) -> Node {
        match node {
            parsed::Node::Error => self.node(NodeKind::Error, Type::Unit),
            parsed::Node::Id(id) => {
                let Some(sym) = self.syms.get(*id) else {
                    self.error("undefined >~<");
                    return self.node(NodeKind::Error, Type::Error);
                };
                self.node(NodeKind::Id(*id), sym.typ.clone())
            }
            parsed::Node::Int(value) => self.node(
                NodeKind::Int(*value),
                if *value > i32::MAX as i64 {
                    Type::U32
                } else {
                    Type::I32
                },
            ),
            parsed::Node::String(value) => self.node(NodeKind::String(value.clone()), Type::String),
            parsed::Node::Group(expr) => {
                let expr = self.check_expr(expr);
                let typ = expr.typ.clone();
                self.node(NodeKind::Group(Box::new(expr)), typ)
            }
            parsed::Node::Block(stmts) => {
                self.syms.enter_scope();
                let stmts = stmts
                    .iter()
                    .map(|stmt| self.check_expr(stmt))
                    .collect::<Vec<_>>();
                self.syms.leave_scope().unwrap();
                let typ = stmts
                    .last()
                    .map(|stmt| stmt.typ.clone())
                    .unwrap_or(Type::Unit);
                self.node(NodeKind::Block(stmts), typ)
            }
            parsed::Node::Call { subject, args } => {
                let subject = Box::new(self.check_expr(subject));

                let args = args
                    .iter()
                    .map(|arg| self.check_expr(arg))
                    .collect::<Vec<_>>();
                let typ = 'br: {
                    match subject.typ.clone() {
                        Type::Fn {
                            id: _,
                            params,
                            return_typ,
                        } => {
                            if args.len() != params.len() {
                                self.error("too few/many args");
                                break 'br Type::Error;
                            }
                            if args
                                .iter()
                                .zip(params)
                                .map(|(arg, param)| self.compatible(&arg.typ, &param.typ))
                                .any(|is_compatible| !is_compatible)
                            {
                                self.error("incorrect args");
                                break 'br Type::Error;
                            }
                            *return_typ
                        }
                        _ => {
                            self.error("not a function");
                            Type::Error
                        }
                    }
                };
                self.node(NodeKind::Call { subject, args }, typ)
            }
            parsed::Node::If {
                cond,
                truthy,
                falsy,
            } => {
                let cond = Box::new(self.check_expr(cond));
                let truthy = Box::new(self.check_expr(truthy));
                let falsy = falsy.as_ref().map(|block| Box::new(self.check_expr(block)));
                let typ = 'br: {
                    match falsy.as_ref().map(|block| block.typ.clone()) {
                        Some(falsy_typ) => {
                            if !self.compatible(&truthy.typ, &falsy_typ) {
                                self.error("incompatible types #2");
                                break 'br Type::Error;
                            }
                            falsy_typ
                        }
                        None => Type::Unit,
                    }
                };
                self.node(
                    NodeKind::If {
                        cond,
                        truthy,
                        falsy,
                    },
                    typ,
                )
            }
            parsed::Node::Loop { body } => {
                self.syms.enter_scope();
                let body = Box::new(self.check_expr(body));
                let typ = body.typ.clone();
                self.node(NodeKind::Loop { body }, typ)
            }
            parsed::Node::Break => self.node(NodeKind::Break, Type::Unit),
            parsed::Node::Assign { subject, value } => {
                let subject = Box::new(self.check_expr(subject));
                let value = Box::new(self.check_expr(value));

                let typ = if !self.compatible(&subject.typ, &value.typ) {
                    self.error("incompatible types #3");
                    Type::Error
                } else {
                    subject.typ.clone()
                };
                self.node(NodeKind::Assign { subject, value }, typ)
            }
            parsed::Node::Let { subject, value } => {
                let (subject, subject_typ) = match subject.as_ref() {
                    parsed::Node::Param { subject, typ } => {
                        (subject, typ.as_ref().map(|typ| self.check_type(typ)))
                    }
                    _ => unreachable!(),
                };

                let value = Box::new(self.check_expr(value));
                let typ = value.typ.clone();

                if subject_typ
                    .as_ref()
                    .map(|subject_typ| !self.compatible(subject_typ, &typ))
                    .unwrap_or(false)
                {
                    self.error("incompatible types #1");
                    return self.node(NodeKind::Error, Type::Error);
                }

                let subject = match subject.as_ref() {
                    parsed::Node::Id(id) => {
                        if self.syms.defined_locally(*id) {
                            self.error("redefinition");
                            return self.node(NodeKind::Error, Type::Error);
                        }
                        self.syms.define(*id, typ.clone());
                        Box::new(self.node(
                            NodeKind::Param {
                                subject: Box::new(self.node(NodeKind::Id(*id), Type::Unit)),
                                typ: Some(Type::Unit),
                            },
                            Type::Unit,
                        ))
                    }
                    _ => unreachable!(),
                };

                self.node(NodeKind::Let { subject, value }, typ)
            }
            parsed::Node::Fn {
                subject,
                params: _,
                return_typ: _,
                body,
            } => {
                let parsed::Node::Id(id) = subject.as_ref() else { unreachable!() };

                let Some(sym) = self.syms.get(*id).cloned() else {
                    // rejected in fn scanner
                    return self.node(NodeKind::Error,Type::Error);
                };

                let Type::Fn { id: fn_id, params, return_typ } = sym.typ else {
                    self.error("redefintion");
                    return self.node(NodeKind::Error,Type::Error);
                };

                self.syms.enter_scope();

                for param in &params {
                    let NodeKind::Param {
                        ref subject,
                        typ: Some(ref typ),
                    } = param.kind else { unreachable!() };
                    let NodeKind::Id(id) = subject.kind else { unreachable!() };
                    self.syms.define(id, typ.clone());
                }

                let body = Box::new(self.check_expr(body));

                if !self.compatible(&return_typ, &body.typ) {
                    self.error("return type violated");
                }

                self.syms.leave_scope().unwrap();

                self.node(
                    NodeKind::Fn {
                        subject: Box::new(self.node(NodeKind::Id(*id), Type::Unit)),
                        params,
                        return_typ: *return_typ,
                        body,
                    },
                    Type::Unit,
                )
            }
            parsed::Node::Return { value } => {
                let value = value.as_ref().map(|value| Box::new(self.check_expr(value)));
                let typ = value
                    .as_ref()
                    .map(|value| value.typ.clone())
                    .unwrap_or(Type::Unit);
                self.node(NodeKind::Return { value }, typ)
            }
            parsed::Node::Param { .. } => unreachable!("handle elsewhere"),
        }
    }

    fn check_type(&self, node: &parsed::Node) -> Type {
        match node {
            parsed::Node::Error => Type::Error,
            parsed::Node::Id(value) => {
                if *value == hash("i32") {
                    Type::I32
                } else if *value == hash("u32") {
                    Type::U32
                } else {
                    todo!("symbol lookup idk")
                }
            }
            _ => unreachable!(),
        }
    }

    fn compatible(&self, typ_a: &Type, typ_b: &Type) -> bool {
        typ_a == typ_b
    }

    fn node(&self, kind: NodeKind, typ: Type) -> Node {
        Node {
            kind,
            typ,
            table_id: self.syms.table_id(),
        }
    }

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

#[test]
fn test_checker() {
    use crate::parser::Parser;
    use pretty_assertions::assert_eq;
    use NodeKind::{Block, Call, Error, Fn, Id, Int, Let, Param, Return};
    use Type::{Unit, I32};

    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 check = |text| Checker::<SeqIdGen>::new_with_fn_id_gen().check(&Parser::new(text).parse());

    assert_eq!(
        check("let a = 5; a;"),
        vec![
            Node {
                kind: Let {
                    subject: Box::new(Node {
                        kind: Param {
                            subject: Box::new(Node {
                                kind: Id(hash("a")),
                                typ: Unit,
                                table_id: 0,
                            }),
                            typ: Some(Unit)
                        },
                        typ: Unit,
                        table_id: 0,
                    }),
                    value: Box::new(Node {
                        kind: Int(5),
                        typ: I32,
                        table_id: 0,
                    })
                },
                typ: I32,
                table_id: 0,
            },
            Node {
                kind: Id(hash("a")),
                typ: I32,
                table_id: 0,
            }
        ]
    );

    assert_eq!(
        check("let a = 5; { a; }"),
        vec![
            Node {
                kind: Let {
                    subject: Box::new(Node {
                        kind: Param {
                            subject: Box::new(Node {
                                kind: Id(hash("a")),
                                typ: Unit,
                                table_id: 0,
                            }),
                            typ: Some(Unit)
                        },
                        typ: Unit,
                        table_id: 0,
                    }),
                    value: Box::new(Node {
                        kind: Int(5),
                        typ: I32,
                        table_id: 0,
                    })
                },
                typ: I32,
                table_id: 0,
            },
            Node {
                kind: Block(vec![Node {
                    kind: Id(hash("a")),
                    typ: I32,
                    table_id: 1,
                },]),
                typ: I32,
                table_id: 0,
            },
        ]
    );

    assert_eq!(
        check("let a = 5; a; { a; let b = 5; b; } a; b;"),
        vec![
            Node {
                kind: Let {
                    subject: Box::new(Node {
                        kind: Param {
                            subject: Box::new(Node {
                                kind: Id(hash("a")),
                                typ: Unit,
                                table_id: 0,
                            }),
                            typ: Some(Unit)
                        },
                        typ: Unit,
                        table_id: 0,
                    }),
                    value: Box::new(Node {
                        kind: Int(5),
                        typ: I32,
                        table_id: 0,
                    })
                },
                typ: I32,
                table_id: 0,
            },
            Node {
                kind: Id(hash("a")),
                typ: I32,
                table_id: 0,
            },
            Node {
                kind: Block(vec![
                    Node {
                        kind: Id(hash("a")),
                        typ: I32,
                        table_id: 1,
                    },
                    Node {
                        kind: Let {
                            subject: Box::new(Node {
                                kind: Param {
                                    subject: Box::new(Node {
                                        kind: Id(hash("b")),
                                        typ: Unit,
                                        table_id: 1,
                                    }),
                                    typ: Some(Unit)
                                },
                                typ: Unit,
                                table_id: 1,
                            }),
                            value: Box::new(Node {
                                kind: Int(5),
                                typ: I32,
                                table_id: 1,
                            })
                        },
                        typ: I32,
                        table_id: 1,
                    },
                    Node {
                        kind: Id(hash("b")),
                        typ: I32,
                        table_id: 1,
                    }
                ]),
                typ: I32,
                table_id: 0,
            },
            Node {
                kind: Id(hash("a")),
                typ: I32,
                table_id: 0,
            },
            Node {
                kind: Error,
                typ: Type::Error,
                table_id: 0,
            }
        ]
    );

    assert_eq!(
        check("fn a() -> i32 { return 0; } \n a();"),
        vec![
            Node {
                kind: Fn {
                    subject: Box::new(Node {
                        kind: Id(hash("a")),
                        typ: Unit,
                        table_id: 0,
                    }),
                    params: vec![],
                    return_typ: I32,
                    body: Box::new(Node {
                        kind: Block(vec![Node {
                            kind: Return {
                                value: Some(Box::new(Node {
                                    kind: Int(0),
                                    typ: I32,
                                    table_id: 2,
                                }),),
                            },
                            typ: I32,
                            table_id: 2,
                        },],),
                        typ: I32,
                        table_id: 1,
                    }),
                },
                typ: Unit,
                table_id: 0,
            },
            Node {
                kind: Call {
                    subject: Box::new(Node {
                        kind: Id(hash("a")),
                        typ: Type::Fn {
                            id: 0,
                            params: vec![],
                            return_typ: Box::new(I32)
                        },
                        table_id: 0
                    }),
                    args: vec![]
                },
                typ: I32,
                table_id: 0,
            },
        ]
    );
}