use std::cell::Cell;
use std::collections::{BTreeMap, HashMap};

use super::parser;
use super::parser::block_expression::BlockExpr;
pub use callable::{Callable, CallableContext, CallableId};
pub use types::{Type, TypeStruct, TypingContext};

mod callable;
pub mod lowering;
mod pretty_ir;
mod type_infer;
pub mod typed_ir;
pub mod types;

#[cfg(never)]
use crate::builtin_cells::get_builtins;

// pub use lowering::lower_module;

fn make_pubid(id: &str) -> String {
    "\\".to_owned() + id
}

#[derive(Debug)]
pub enum CompileErrorKind {
    UndefinedReference(String),
    BadArgCount { received: usize, expected: usize },
    TodoError(String),
    TypeError { expected: Type, found: Type },
}

#[derive(Debug)]
pub struct CompileError {
    kind: CompileErrorKind,
}

impl CompileError {
    fn new(kind: CompileErrorKind) -> Self {
        Self { kind }
    }
}

#[derive(Debug)]
pub struct Module {
    pub blocks: Vec<typed_ir::Body>,
}

pub struct Context {
    /// map callable name to callable
    callable_names: BTreeMap<String, CallableId>,
    callables: CallableContext,
    /// type names
    typenames: BTreeMap<String, Type>,
    types: TypingContext,
    /// map signal name to Signal
    signals: BTreeMap<String, typed_ir::Signal>,
    /// incrementing counter for unique IDs
    ids: Counter,
}

struct Counter(Cell<usize>);

impl Counter {
    fn new() -> Counter {
        Counter(Cell::new(0))
    }
    fn next(&self) -> usize {
        let next = self.0.get() + 1;
        self.0.set(next);
        next
    }
}

impl Context {
    pub fn new() -> Self {
        let mut tcx = TypingContext::new();
        let ccx = CallableContext::new(&mut tcx);

        let typenames = [
            ("Logic".to_string(), tcx.primitives.logic),
            ("Num".to_string(), tcx.primitives.elabnum),
        ]
        .into();

        let callable_names = [("reduce_or".to_string(), ccx.builtins.reduce_or)].into();

        Context {
            callables: ccx,
            callable_names,
            signals: BTreeMap::new(),
            types: tcx,
            typenames,
            ids: Counter::new(),
        }
    }

    fn try_get_signal(&self, signame: &str) -> Result<&typed_ir::Signal, CompileError> {
        self.signals.get(signame).ok_or_else(|| {
            CompileError::new(CompileErrorKind::UndefinedReference(signame.to_owned()))
        })
    }

    fn try_get_type(&self, typename: &str) -> Result<Type, CompileError> {
        self.typenames.get(typename).copied().ok_or_else(|| {
            CompileError::new(CompileErrorKind::UndefinedReference(typename.to_owned()))
        })
    }

    fn try_get_callable(&self, callname: &str) -> Result<CallableId, CompileError> {
        self.callable_names.get(callname).copied().ok_or_else(|| {
            CompileError::new(CompileErrorKind::UndefinedReference(callname.to_owned()))
        })
    }

    fn eval_expression(&self, expr: &typed_ir::Expr) -> Result<types::ElabData, CompileError> {
        match &expr.kind {
            typed_ir::ExprKind::Literal(lit) => Ok(lit.clone()),
            typed_ir::ExprKind::Path(_) => todo!("evaluate path"),
            typed_ir::ExprKind::Call(_call) => todo!("evaluate call"),
            typed_ir::ExprKind::Match(_) => todo!(),
        }
    }

    fn intern_expression(
        &self,
        exprs: &mut HashMap<typed_ir::ExprId, typed_ir::Expr>,
        expr: typed_ir::Expr,
    ) -> typed_ir::ExprId {
        let expr_id = expr.id;
        exprs.insert(expr.id, expr);
        expr_id
    }

    fn type_expression(
        &mut self,
        exprs: &mut HashMap<typed_ir::ExprId, typed_ir::Expr>,
        expr: &parser::expression::Expression,
    ) -> Result<typed_ir::ExprId, CompileError> {
        use parser::expression::Expression;
        let id = typed_ir::ExprId(self.ids.next() as u32);
        let t_expr = match expr {
            Expression::Path(name) => {
                let signal = self.try_get_signal(name)?;
                let this_expr = typed_ir::Expr {
                    id,
                    kind: typed_ir::ExprKind::Path(signal.id),
                    typ: signal.typ,
                };
                self.intern_expression(exprs, this_expr)
            }
            Expression::Literal(lit) => {
                let data = match lit.kind {
                    parser::expression::LiteralKind::Num(num) => self.types.make_elabnum_u32(num),
                    parser::expression::LiteralKind::Const(width, val) => {
                        self.types.make_const_u32(width, val as u32)
                    }
                };
                let this_expr = typed_ir::Expr {
                    id,
                    kind: typed_ir::ExprKind::Literal(data),
                    typ: self.types.primitives.infer,
                };
                self.intern_expression(exprs, this_expr)
            }
            Expression::UnOp(op) => {
                let a = self.type_expression(exprs, &op.a)?;
                let this_expr = typed_ir::Expr {
                    id,
                    kind: typed_ir::ExprKind::Call(typed_ir::Call {
                        called: self.callables.builtins.bitnot,
                        args: vec![a],
                        genargs: vec![],
                    }),
                    typ: self.types.primitives.infer,
                };
                self.intern_expression(exprs, this_expr)
            }
            Expression::BinOp(op) => {
                let (a, b) = (
                    self.type_expression(exprs, &op.a)?,
                    self.type_expression(exprs, &op.b)?,
                );
                let this_expr = typed_ir::Expr {
                    id,
                    kind: typed_ir::ExprKind::Call(typed_ir::Call {
                        called: self.callables.builtins.xor,
                        args: vec![a, b],
                        genargs: vec![],
                    }),
                    typ: self.types.primitives.infer,
                };
                self.intern_expression(exprs, this_expr)
            }
            Expression::Call(call) => {
                let args_resolved = call
                    .args
                    .iter()
                    .map(|expr| self.type_expression(exprs, expr))
                    .collect::<Result<Vec<_>, _>>()?;
                let called = self.try_get_callable(call.name.fragment())?;
                let called_callable = self.callables.get(called);
                if args_resolved.len() != called_callable.argcount() {
                    return Err(CompileError::new(CompileErrorKind::BadArgCount {
                        received: args_resolved.len(),
                        expected: called_callable.argcount(),
                    }));
                };
                let genargs_resolved = called_callable
                    .genargs
                    .iter()
                    .map(|genarg| genarg.1)
                    .collect();
                let this_expr = typed_ir::Expr {
                    id,
                    kind: typed_ir::ExprKind::Call(typed_ir::Call {
                        called,
                        args: args_resolved,
                        genargs: genargs_resolved,
                    }),
                    typ: self.types.primitives.infer,
                };
                self.intern_expression(exprs, this_expr)
            }
            Expression::BlockExpr(block) => match &**block {
                BlockExpr::IfElse(_) => todo!(),
                BlockExpr::Match(match_) => {
                    let expr = self.type_expression(exprs, &match_.expr)?;
                    let arms = match_
                        .arms
                        .iter()
                        .map(|(cond, val)| {
                            Ok((
                                self.type_expression(exprs, cond)?,
                                self.type_expression(exprs, val)?,
                            ))
                        })
                        .collect::<Result<_, _>>()?;
                    let typed = typed_ir::Match { expr, arms };
                    let this_expr = typed_ir::Expr {
                        id,
                        kind: typed_ir::ExprKind::Match(Box::new(typed)),
                        typ: self.types.primitives.infer,
                    };
                    self.intern_expression(exprs, this_expr)
                }
                BlockExpr::Block(block) => {
                    // TODO: we need to find some way of resolving a name to an expression
                    todo!("can not convert blocks to typed_ir yet");
                    for (kind, name, expr) in &block.assignments {
                        let signal = typed_ir::Signal {
                            id: typed_ir::DefId(self.ids.next() as u32),
                            typ: self.types.primitives.infer,
                        };
                        // TODO: need to add this signal to the block from here, somehow
                        self.signals.insert(name.span().to_string(), signal);
                    }
                    self.type_expression(exprs, &block.value)?
                }
            },
            Expression::StructInit(_) => todo!("structure initialization"),
        };
        Ok(t_expr)
    }

    fn callable_from_block(
        &mut self,
        comb: &parser::comb::CombBlock,
    ) -> Result<CallableId, CompileError> {
        let mut genargs = vec![];
        for (idx, tvarname) in comb.genparams.iter().enumerate() {
            let tvar = self.types.make_typevar(0, idx as u32);
            // hacky workaround for no scopes
            self.typenames.insert(tvarname.name.to_string(), tvar);
            genargs.push((Some(tvarname.name.to_string()), tvar));
        }

        let mut args = vec![];
        for port in comb.ports.iter() {
            let sig_typename = &port.net.typ;
            let mut sig_type = self.try_get_type(sig_typename.name.fragment())?;
            if let Some(arg) = &sig_typename.generics {
                let mut exprs = Default::default();
                let elab_expr = self.type_expression(&mut exprs, arg)?;
                let elab_val = self.eval_expression(exprs.get(&elab_expr).unwrap())?;
                sig_type = self
                    .types
                    .parameterize(sig_type, &[types::GenericArg::Elab(elab_val)])
                    .unwrap();
            }

            args.push((Some(port.net.name.to_string()), sig_type));
        }

        let ret_typename = &comb.ret.name;
        let ret_type = self.try_get_type(ret_typename.fragment())?;

        let signature = Callable {
            name: comb.name.to_string(),
            args,
            genargs,
            ret_type,
        };
        let signature_id = self.callables.add(signature);
        self.callable_names
            .insert(comb.name.to_string(), signature_id);
        Ok(signature_id)
    }

    fn type_comb(
        &mut self,
        comb: &parser::comb::CombBlock,
    ) -> Result<typed_ir::Body, CompileError> {
        let mut signals = Vec::new();

        let callable_id = self.callable_from_block(comb)?;
        let callable = self.callables.get(callable_id);

        for (argname, typ) in &callable.args {
            let sig_id = self.ids.next();
            let sig = typed_ir::Signal {
                id: typed_ir::DefId(sig_id as u32),
                typ: *typ,
            };
            signals.push(sig.clone());
            if let Some(argname) = argname {
                self.signals.insert(argname.clone(), sig);
            }
        }

        let mut exprs = Default::default();

        // TODO: fixme
        let root_expr = self.type_expression(&mut exprs, &comb.expr.value)?;

        Ok(typed_ir::Body {
            signature: callable_id,
            signals,
            exprs,
            expr: root_expr,
        })
    }

    pub fn type_module(&mut self, module: parser::Module) -> Result<Module, CompileError> {
        let mut typed_module = Module { blocks: vec![] };
        // hacky loop to predefine all names
        for item in &module.items {
            match &item {
                parser::ModuleItem::Comb(comb) => self.callable_from_block(comb)?,
                parser::ModuleItem::Proc(_) => todo!("proc block"),
                parser::ModuleItem::State(_) => todo!("state block"),
                parser::ModuleItem::Struct(_) => todo!("struct block"),
            };
        }
        for item in module.items {
            let block = match &item {
                parser::ModuleItem::Comb(comb) => self.type_comb(comb)?,
                parser::ModuleItem::Proc(_) => todo!("proc block"),
                parser::ModuleItem::State(_) => todo!("state block"),
                parser::ModuleItem::Struct(_) => todo!("struct block"),
            };
            typed_module.blocks.push(block);
        }
        Ok(typed_module)
    }
}