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 }
    }
}

/// A user-defined signal
pub struct Signal {
    /// the user-visible name of the signal
    pub name: String,
    /// the id of the signal in RTLIL
    pub il_id: String,
    /// the type of the signal
    pub typ: Type,
    // unique ID of the signal
    // pub uid: u64,
}

impl Signal {
    fn sigspec(&self) -> rtlil::SigSpec {
        rtlil::SigSpec::Wire(self.il_id.to_owned())
    }
}

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 type_comb(
        &mut self,
        comb: &parser::comb::CombBlock,
    ) -> Result<typed_ir::Block, CompileError> {
        let mut signals = Vec::new();

        for (idx, tvarname) in comb.genparams.iter().enumerate() {
            let tvar = self.types.make_typevar(0, idx as u32);
            self.typenames.insert(tvarname.name.to_string(), tvar);
        }

        for port in comb.ports.iter() {
            let sig_id = self.ids.next();

            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();
            }

            let sig = typed_ir::Signal {
                id: typed_ir::DefId(sig_id as u32),
                typ: sig_type,
            };
            signals.push(sig.clone());
            self.signals.insert(port.net.name.to_string(), sig);
        }

        let ret_typename = &comb.ret.name;
        let _ret_type = self.try_get_type(ret_typename.fragment())?;
        // TODO: use ret type

        let root_expr = self.type_expression(&comb.expr)?;

        Ok(typed_ir::Block {
            signals,
            expr: root_expr,
        })
    }

    pub fn type_module(&mut self, module: parser::Module) -> Result<typed_ir::Block, CompileError> {
        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"),
            };
            return Ok(block);
        }
        Err(CompileError::new(CompileErrorKind::TodoError(
            "no blocks in module".to_string(),
        )))
    }

    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(_) => todo!(),
            // 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 {
        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(_) => todo!(),
            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(())
    }
}