use std::collections::BTreeMap;

use super::types::{make_primitives, TypeStruct};
use super::{make_pubid, CompileError, CompileErrorKind, Context, Signal, TODO_WIDTH};
use crate::builtin_cells::get_builtins;
use crate::parser;
use crate::parser::expression::Expression;
use crate::rtlil;
use crate::rtlil::RtlilWrite;

/// context used when generating processes
struct ProcContext {
    updates: Vec<(rtlil::SigSpec, rtlil::SigSpec)>,
    next_sigs: BTreeMap<String, rtlil::SigSpec>,
}

fn lower_process_statement(
    ctx: &Context,
    pctx: &mut ProcContext,
    module: &mut rtlil::Module,
    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_sig;
            if let Some(sig) = pctx.next_sigs.get(assig.lhs) {
                next_sig = sig.clone();
            } else {
                let next_gen_id = format!("${}$next", assig.lhs);
                module.add_wire(rtlil::Wire::new(&next_gen_id, TODO_WIDTH, None));
                next_sig = rtlil::SigSpec::Wire(next_gen_id);

                pctx.next_sigs
                    .insert(assig.lhs.to_owned(), next_sig.clone());

                // trigger the modified value to update
                pctx.updates
                    .push((ctx.try_get_signal(assig.lhs)?.sigspec(), next_sig.clone()));
            };

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

            rtlil::CaseRule {
                assign: vec![(next_sig, 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, pctx, module, &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 pctx = ProcContext {
        updates: vec![],
        next_sigs: BTreeMap::new(),
    };
    let mut cases = vec![];
    for stmt in &process.items {
        let case = lower_process_statement(ctx, &mut pctx, module, stmt)?;
        cases.push(case);
    }

    let sync_sig = ctx.try_get_signal(process.net.fragment())?;
    let sync_cond = rtlil::SyncCond::Posedge(sync_sig.sigspec());
    let sync_rule = rtlil::SyncRule {
        cond: sync_cond,
        assign: pctx.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 desugar_binop<'a>(op: parser::expression::BinOp<'a>) -> parser::expression::Call<'a> {
    let a = desugar_expression(op.a);
    let b = desugar_expression(op.b);
    let op_func = match op.kind {
        parser::expression::BinOpKind::And => "and",
        parser::expression::BinOpKind::Or => "or",
        parser::expression::BinOpKind::Xor => "xor",
    };
    parser::expression::Call {
        name: op_func.into(),
        args: vec![a, b],
    }
}

fn desugar_unop<'a>(op: parser::expression::UnOp<'a>) -> parser::expression::Call<'a> {
    let a = desugar_expression(op.a);
    let op_func = match op.kind {
        parser::expression::UnOpKind::BitNot => "not",
        parser::expression::UnOpKind::Not => todo!(),
    };
    parser::expression::Call {
        name: op_func.into(),
        args: vec![a],
    }
}

fn desugar_expression<'a>(expr: Expression<'a>) -> Expression<'a> {
    // TODO: allow ergonomic traversal of AST
    match expr {
        Expression::Path(_) => expr,
        Expression::Literal(_) => expr,
        Expression::Call(mut call) => {
            let new_args = call.args.into_iter().map(desugar_expression).collect();
            call.args = new_args;
            Expression::Call(call)
        }
        Expression::BinOp(op) => Expression::Call(Box::new(desugar_binop(*op))),
        Expression::UnOp(op) => Expression::Call(Box::new(desugar_unop(*op))),
    }
}

fn lower_expression(
    ctx: &Context,
    module: &mut rtlil::Module,
    expr: &Expression,
) -> Result<rtlil::SigSpec, CompileError> {
    let expr = desugar_expression(expr.clone());
    match expr {
        Expression::Path(ident) => {
            let signal = ctx.try_get_signal(ident)?;
            Ok(signal.sigspec())
        }
        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.argcount() {
                return Err(CompileError::new(CompileErrorKind::BadArgCount {
                    expected: callable.argcount(),
                    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)
        }
        // TODO: instantiate operators directly here instead of desugaring, once the callable infrastructure improves
        // to get better errors
        Expression::Literal(lit) => Ok(rtlil::SigSpec::Const(
            lit.span().fragment().parse().unwrap(),
            TODO_WIDTH,
        )),
        Expression::UnOp(_) => todo!(),
        Expression::BinOp(_) => todo!(),
    }
}

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

fn lower_comb(
    ctx: &mut Context,
    module: &mut rtlil::Module,
    pa_comb: parser::comb::CombBlock,
) -> Result<(), CompileError> {
    for (num, port) in pa_comb.ports.iter().enumerate() {
        let port_id = make_pubid(port.net.name.fragment());
        let port_tyname = &port.net.typ;
        ctx.try_get_type(port_tyname.name.fragment())?;
        module.add_wire(rtlil::Wire::new(
            port_id.clone(),
            TODO_WIDTH,
            Some(rtlil::PortOption::Input((num + 1) as i32)),
        ));
        let typ = TypeStruct::logic_width(TODO_WIDTH);
        let signal = Signal {
            name: port.net.name.fragment().to_string(),
            il_id: port_id,
            // TODO: CRIMES CRIMES CRIMES
            typ: Box::leak(Box::new(typ)),
        };
        ctx.signals
            .insert(port.net.name.fragment().to_string(), signal);
    }
    let ret_id = module.make_genid("ret");
    module.add_wire(rtlil::Wire::new(
        ret_id.clone(),
        TODO_WIDTH,
        Some(rtlil::PortOption::Output(99)),
    ));
    let out_sig = lower_expression(ctx, module, &pa_comb.expr)?;
    module.add_connection(&rtlil::SigSpec::Wire(ret_id), &out_sig);
    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("test"));

    writer.write_line("autoidx 1");
    ir_module.write_rtlil(&mut writer);
    Ok(writer.finish())
}