use std::cell::Cell;
use std::collections::BTreeMap;
use std::fmt::Write;

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

mod callable;
pub mod lowering;
pub mod typed_ir;
pub mod types;

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

// pub use lowering::lower_module;

/// lots of code is still not width-aware, this constant keeps track of that
const TODO_WIDTH: u32 = 1;

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::Block>,
}

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 {
                called,
                args,
                genargs,
            } => todo!("evaluate call"),
        }
    }

    fn type_expression(
        &self,
        expr: &parser::expression::Expression,
    ) -> Result<typed_ir::Expr, 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)?;
                typed_ir::Expr {
                    id,
                    kind: typed_ir::ExprKind::Path(signal.id),
                    typ: signal.typ,
                }
            }
            Expression::Literal(lit) => {
                // TODO: make this a proper enum instead of having to match on everything
                let data = match lit.kind() {
                    parser::tokens::TokenKind::Number => {
                        let num = lit.span().fragment().parse().unwrap();
                        self.types.make_elabnum_u32(num)
                    }
                    _ => unreachable!("non-literal token in literal?"),
                };
                typed_ir::Expr {
                    id,
                    kind: typed_ir::ExprKind::Literal(data),
                    typ: self.types.primitives.infer,
                }
            }
            Expression::UnOp(op) => {
                let a = self.type_expression(&op.a)?;
                typed_ir::Expr {
                    id,
                    kind: typed_ir::ExprKind::Call {
                        called: self.callables.builtins.bitnot,
                        args: vec![a],
                        genargs: vec![],
                    },
                    typ: self.types.primitives.infer,
                }
            }
            Expression::BinOp(op) => {
                let (a, b) = (self.type_expression(&op.a)?, self.type_expression(&op.b)?);
                typed_ir::Expr {
                    id,
                    kind: typed_ir::ExprKind::Call {
                        called: self.callables.builtins.xor,
                        args: vec![a, b],
                        genargs: vec![],
                    },
                    typ: self.types.primitives.infer,
                }
            }
            Expression::Call(call) => {
                let args_resolved = call
                    .args
                    .iter()
                    .map(|expr| self.type_expression(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();
                typed_ir::Expr {
                    id,
                    kind: typed_ir::ExprKind::Call {
                        called,
                        args: args_resolved,
                        genargs: genargs_resolved,
                    },
                    typ: self.types.primitives.infer,
                }
            }
        };
        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 elab_expr = self.type_expression(arg)?;
                let elab_val = self.eval_expression(&elab_expr)?;
                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::Block, 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 root_expr = self.type_expression(&comb.expr)?;

        Ok(typed_ir::Block {
            signature: callable_id,
            signals,
            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"),
            };
        }
        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"),
            };
            typed_module.blocks.push(block);
        }
        Ok(typed_module)
    }

    pub fn infer_expr_types(&mut self, expr: &typed_ir::Expr) -> typed_ir::Expr {
        if self.types.is_fully_typed(expr.typ) {
            // there is nothing more to infer
            return expr.clone();
        }
        match &expr.kind {
            typed_ir::ExprKind::Literal(lit) => expr.clone().with_type(lit.typ),
            // we can not see beyond this expression right now
            typed_ir::ExprKind::Path(_) => expr.clone(),
            typed_ir::ExprKind::Call {
                called,
                args,
                genargs,
            } => {
                let args_typed: Vec<_> = args.iter().map(|ex| self.infer_expr_types(ex)).collect();
                let callee_def = self.callables.get(*called);

                let param_types: Vec<_> = callee_def.args.iter().map(|param| param.1).collect();
                let inferred_args: Vec<_> = param_types
                    .iter()
                    .zip(&args_typed)
                    .map(|(param, arg)| self.types.infer_type(*param, arg.typ))
                    .collect();

                let mut genargs: Vec<_> = callee_def.genargs.iter().map(|a| a.1).collect();

                let mut new_type = callee_def.ret_type;

                if genargs.len() != 0 {
                    // need to infer generic arguments
                    for inf_res in inferred_args {
                        match inf_res {
                            types::InferenceResult::First(_) => todo!(),
                            types::InferenceResult::Second(_) => todo!(),
                            types::InferenceResult::TypeVar(dbi, tvar, typ) => {
                                assert_eq!(dbi, 0);
                                // TODO: type check argument instead of just using it
                                genargs[tvar as usize] = typ;
                            }
                            types::InferenceResult::Incompatible => todo!(),
                            types::InferenceResult::Ambigous => todo!(),
                        }
                    }

                    // TODO: HACKY HACKY HACK
                    new_type = genargs[0];
                }

                let mut new_expr = expr.clone();
                new_expr.typ = new_type;
                new_expr.kind = typed_ir::ExprKind::Call {
                    called: called.clone(),
                    args: args_typed,
                    genargs,
                };
                new_expr
            }
        }
    }

    pub fn infer_types(&mut self, mut block: typed_ir::Block) -> typed_ir::Block {
        let new_root = self.infer_expr_types(&block.expr);
        block.expr = new_root;
        block
    }

    pub fn pretty_typed_block(
        &self,
        w: &mut dyn std::fmt::Write,
        block: &typed_ir::Block,
    ) -> std::fmt::Result {
        let callsig = self.callables.get(block.signature);
        {
            // TODO: ugly copy paste job
            let args = callsig
                .args
                .iter()
                .map(|(name, typ)| {
                    let mut out = String::new();
                    self.types.pretty_type(&mut out, *typ)?;
                    Ok(out)
                })
                .collect::<Result<Vec<String>, std::fmt::Error>>()?;
            let genargs = callsig
                .genargs
                .iter()
                .map(|(name, typ)| {
                    let mut type_str = String::new();
                    self.types.pretty_type(&mut type_str, *typ)?;
                    Ok(type_str)
                })
                .collect::<Result<Vec<String>, std::fmt::Error>>()?;
            writeln!(
                w,
                "block {}<{}>({})",
                callsig.name(),
                genargs.join(", "),
                args.join(", ")
            )?;
        }
        for sig in &block.signals {
            let mut typ_pretty = String::new();
            self.types.pretty_type(&mut typ_pretty, sig.typ)?;
            writeln!(w, "sig_{}: {}", sig.id.0, typ_pretty)?
        }
        self.pretty_typed_expr(w, &block.expr)?;
        Ok(())
    }

    pub fn pretty_typed_expr(
        &self,
        w: &mut dyn std::fmt::Write,
        expr: &typed_ir::Expr,
    ) -> std::fmt::Result {
        let expr_pretty = match &expr.kind {
            typed_ir::ExprKind::Literal(lit) => {
                let mut lit_str = String::new();
                self.types.pretty_value(&mut lit_str, lit)?;
                lit_str
            }
            typed_ir::ExprKind::Path(path) => format!("sig_{}", path.0),
            typed_ir::ExprKind::Call {
                called,
                args,
                genargs,
            } => {
                let args = args
                    .iter()
                    .map(|arg| {
                        self.pretty_typed_expr(w, arg)?;
                        Ok(format!("_{}", arg.id.0))
                    })
                    .collect::<Result<Vec<_>, std::fmt::Error>>()?;
                let callable = self.callables.get(*called);
                let genargs = genargs
                    .iter()
                    .map(|param| {
                        let mut type_str = String::new();
                        self.types.pretty_type(&mut type_str, *param)?;
                        Ok(type_str)
                    })
                    .collect::<Result<Vec<_>, std::fmt::Error>>()?;
                format!(
                    "{}<{}>({})",
                    callable.name(),
                    genargs.join(", "),
                    args.join(", ")
                )
            }
        };
        let mut type_pretty = String::new();
        self.types.pretty_type(&mut type_pretty, expr.typ)?;
        writeln!(w, "let _{}: {} = {}", expr.id.0, type_pretty, expr_pretty)?;
        Ok(())
    }
}