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

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

mod callable;
#[cfg(never)]
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: BTreeMap<CallableId, Callable>,
    /// 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 tcx = TypingContext::new();
        Context {
            callables: BTreeMap::new(),
            callable_names: BTreeMap::new(),
            signals: BTreeMap::new(),
            types: TypingContext::new(),
            typenames: [
                ("Logic".to_string(), tcx.primitives.logic),
                ("Num".to_string(), tcx.primitives.elabnum),
            ]
            .into(),
            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 } => 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: typed_ir::DefId(99),
                        args: vec![a],
                    },
                    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: typed_ir::DefId(99),
                        args: vec![a, b],
                    },
                    typ: self.types.primitives.infer,
                }
            }
            Expression::Call(call) => {
                let args_resolved = call
                    .args
                    .iter()
                    .map(|expr| self.type_expression(expr))
                    .collect::<Result<Vec<_>, _>>()?;
                typed_ir::Expr {
                    id,
                    kind: typed_ir::ExprKind::Call {
                        called: typed_ir::DefId(99),
                        args: args_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 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 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 } => {
                let args = args
                    .iter()
                    .map(|arg| {
                        self.pretty_typed_expr(w, arg)?;
                        Ok(format!("_{}", arg.id.0))
                    })
                    .collect::<Result<Vec<_>, std::fmt::Error>>()?;
                format!("_{}({})", called.0, 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(())
    }
}