use std::collections::BTreeMap;

use crate::builtin_cells::get_builtins;
use crate::parser;
use crate::rtlil;
use crate::rtlil::RtlilWrite;

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

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

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

pub enum GenericParam<T> {
    Unsolved,
    Solved(T),
}

pub enum Type {
    /// a wire of some width
    Wire(GenericParam<u32>),
}

impl Type {
    pub fn wire() -> Self {
        Self::Wire(GenericParam::Unsolved)
    }
}

pub struct CallArgument {
    pub name: String,
    pub atype: Type,
}

// module that can be instantiated like a function
pub struct Callable {
    pub name: String,
    pub args: Vec<CallArgument>,
    pub ret: Type,
    pub instantiate: Box<dyn Fn(&str, &[rtlil::SigSpec], &rtlil::SigSpec) -> rtlil::Cell>,
}

struct Context {
    callables: BTreeMap<String, Callable>,
}

fn lower_process_statement(
    ctx: &Context,
    module: &mut rtlil::Module,
    updates: &mut Vec<(rtlil::SigSpec, rtlil::SigSpec)>,
    stmt: &parser::proc::ProcStatement,
) -> Result<rtlil::CaseRule, CompileError> {
    let rule = match stmt {
        parser::proc::ProcStatement::IfElse(_) => todo!("if/else unimplemented"),
        parser::proc::ProcStatement::Assign(assig) => {
            // FIXME: actually store this
            let next_gen_id = format!("${}$next", assig.lhs);
            module.add_wire(rtlil::Wire::new(&next_gen_id, TODO_WIDTH, None));

            let next_wire = rtlil::SigSpec::Wire(next_gen_id.clone());
            updates.push((
                rtlil::SigSpec::Wire(assig.lhs.to_owned()),
                next_wire.clone(),
            ));

            let next_expr_wire = lower_expression(ctx, module, &assig.expr)?;

            rtlil::CaseRule {
                assign: vec![(next_wire, next_expr_wire)],
                switches: vec![],
            }
        }
        parser::proc::ProcStatement::Match(match_block) => {
            let match_sig = lower_expression(ctx, module, &match_block.expr)?;
            let mut cases = vec![];
            for arm in &match_block.arms {
                let case = lower_process_statement(ctx, module, updates, &arm.1)?;
                let compare_sig = lower_expression(ctx, module, &arm.0)?;
                cases.push((compare_sig, case));
            }
            let switch_rule = rtlil::SwitchRule {
                signal: match_sig,
                cases,
            };
            rtlil::CaseRule {
                assign: vec![],
                switches: vec![switch_rule],
            }
        }
        parser::proc::ProcStatement::Block(_) => todo!("blocks unimplemented"),
    };
    Ok(rule)
}

fn lower_process(
    ctx: &Context,
    module: &mut rtlil::Module,
    process: &parser::proc::ProcBlock,
) -> Result<(), CompileError> {
    let mut updates = vec![];
    let mut cases = vec![];
    for stmt in &process.items {
        let case = lower_process_statement(ctx, module, &mut updates, &stmt)?;
        cases.push(case);
    }

    let sync_cond = rtlil::SyncCond::Posedge((*process.net.fragment()).into());
    let sync_rule = rtlil::SyncRule {
        cond: sync_cond,
        assign: updates,
    };
    if cases.len() != 1 {
        panic!("only one expression per block, for now")
    }
    assert_eq!(cases.len(), 1);
    let ir_proc = rtlil::Process {
        id: module.make_genid("proc"),
        root_case: cases.into_iter().next().unwrap(),
        sync_rules: vec![sync_rule],
    };
    module.add_process(ir_proc);
    Ok(())
}

fn lower_expression(
    ctx: &Context,
    module: &mut rtlil::Module,
    expr: &parser::Expression,
) -> Result<rtlil::SigSpec, CompileError> {
    match expr {
        parser::Expression::Ident(ident) => Ok(rtlil::SigSpec::Wire(make_pubid(ident))),
        parser::Expression::Call(call) => {
            let args_resolved = call
                .args
                .iter()
                .map(|expr| lower_expression(ctx, module, expr))
                .collect::<Result<Vec<_>, _>>()?;

            let callable = ctx
                .callables
                .get(call.name.fragment() as &str)
                .ok_or_else(|| {
                    CompileError::new(CompileErrorKind::UndefinedReference(
                        call.name.fragment().to_string(),
                    ))
                })?;

            if args_resolved.len() != callable.args.len() {
                return Err(CompileError::new(CompileErrorKind::BadArgCount {
                    expected: callable.args.len(),
                    received: args_resolved.len(),
                }));
            }

            let cell_id = module.make_genid(&callable.name);

            let output_gen_id = format!("{}$out", &cell_id);
            module.add_wire(rtlil::Wire::new(&output_gen_id, TODO_WIDTH, None));
            let output_gen_wire = rtlil::SigSpec::Wire(output_gen_id);

            let cell =
                (*callable.instantiate)(&cell_id, args_resolved.as_slice(), &output_gen_wire);
            module.add_cell(cell);
            Ok(output_gen_wire)
        }
        // operations should really just desugar to callables
        parser::Expression::Operation(_op) => todo!("operators not yet implemented"),
        parser::Expression::Literal(lit) => Ok(rtlil::SigSpec::Const(*lit as i64, TODO_WIDTH)),
    }
}

fn lower_assignment(
    ctx: &Context,
    module: &mut rtlil::Module,
    assignment: parser::Assign,
) -> Result<(), CompileError> {
    let target_id = rtlil::SigSpec::Wire(make_pubid(assignment.lhs));
    let return_wire = lower_expression(ctx, module, &assignment.expr)?;
    module.add_connection(&target_id, &return_wire);
    Ok(())
}

pub fn lower_module(pa_module: parser::Module) -> Result<String, CompileError> {
    let mut writer = rtlil::ILWriter::new();
    let mut ir_module = rtlil::Module::new(make_pubid(pa_module.name));
    let context = Context {
        callables: get_builtins()
            .into_iter()
            .map(|clb| (clb.name.to_owned(), clb))
            .collect(),
    };
    writer.write_line("autoidx 1");
    for (idx, port) in pa_module.ports.iter().enumerate() {
        let dir_option = match port.direction {
            parser::PortDirection::Input => rtlil::PortOption::Input(idx as i32 + 1),
            parser::PortDirection::Output => rtlil::PortOption::Output(idx as i32 + 1),
        };
        let wire = rtlil::Wire::new(
            make_pubid(port.net.name),
            port.net.width.unwrap_or(1) as u32,
            Some(dir_option),
        );
        ir_module.add_wire(wire);
    }
    for item in pa_module.items {
        match item {
            parser::ModuleItem::Assign(assignment) => {
                lower_assignment(&context, &mut ir_module, assignment)?
            }
            parser::ModuleItem::Proc(proc) => lower_process(&context, &mut ir_module, &proc)?,
        }
    }
    ir_module.write_rtlil(&mut writer);
    Ok(writer.finish())
}