use crate::ast::*;
use std::collections::HashMap;
use wasm_encoder::{Function, Instruction, ValType};

pub struct FunctionBuilder {
    locals: Vec<ValType>,
    local_names: HashMap<String, u32>,
    instructions: Vec<Instruction<'static>>,
}

impl FunctionBuilder {
    pub fn new() -> Self {
        Self {
            locals: Vec::new(),
            local_names: HashMap::new(),
            instructions: Vec::new(),
        }
    }

    pub fn push_instr(&mut self, instr: Instruction<'static>) {
        self.instructions.push(instr);
    }

    pub fn add_local(&mut self, name: &str, ty: ValType) -> u32 {
        let index = self.locals.len() as u32;
        self.locals.push(ty);
        self.local_names.insert(name.to_string(), index);
        index
    }

    pub fn get_local(&self, name: &str) -> u32 {
        *self.local_names.get(name).unwrap()
    }

    pub fn finish(self) -> Function {
        let mut func = Function::new_with_locals_types(self.locals.into_iter());

        for instr in self.instructions.into_iter() {
            func.instruction(&instr);
        }

        func.instruction(&Instruction::End);

        func
    }
}

pub fn codegen_branch(ctx: &mut FunctionBuilder, ast: &BranchBody) {
    for stmt in ast.statements.iter() {
        match stmt {
            Statement::Let { ident, expr } => {
                codegen_expr(ctx, expr);
                let idx = ctx.add_local(ident.span.slice, ValType::I32);
                ctx.push_instr(Instruction::LocalSet(idx));
            }
        }
    }

    if let Some(expr) = ast.tail_expr.as_ref() {
        codegen_expr(ctx, expr);
    }
}

pub fn codegen_expr(ctx: &mut FunctionBuilder, ast: &Expr) {
    match ast {
        Expr::BinaryOp(op, sides) => {
            codegen_expr(ctx, &sides.0);
            codegen_expr(ctx, &sides.1);

            let instr = match op {
                BinaryOp::Add => Instruction::I32Add,
            };

            ctx.push_instr(instr);
        }
        Expr::Literal(lit) => match lit {
            Literal::DecimalInteger(span) => {
                let int: i32 = span.slice.parse().unwrap();
                let instr = Instruction::I32Const(int);
                ctx.push_instr(instr);
            }
        },
        Expr::Local(span) => {
            let idx = ctx.get_local(span.slice);
            let instr = Instruction::LocalGet(idx);
            ctx.push_instr(instr);
        }
    }
}