use crate::{
    backend::{elf, symbol::Symbol},
    io::CompilerMsg,
};
use std::collections::HashMap;

pub struct LinkedProgram {
    pub code: Vec<u8>,
    pub start: Option<Addr>,
}

pub struct UnlinkedProgram<I: Instr> {
    pub fns: Vec<UnlinkedFunction<I>>,
    pub ro_data: Vec<(Vec<u8>, Symbol)>,
    pub sym_count: usize,
    pub start: Option<Symbol>,
}

pub struct UnlinkedFunction<I: Instr> {
    pub instrs: Vec<I>,
    pub sym: Symbol,
    // locations within the function for loops & such / goto
    pub locations: HashMap<usize, Symbol>,
}

impl<I: Instr> UnlinkedProgram<I> {
    pub fn link(self) -> Result<LinkedProgram, CompilerMsg> {
        let mut data = Vec::new();
        let mut sym_table = SymTable::new(self.sym_count);
        let mut missing = HashMap::<Symbol, Vec<(I, usize)>>::new();
        for (val, id) in self.ro_data {
            sym_table.insert(id, Addr(data.len() as u64));
            data.extend(val);
        }
        // align
        data.resize(data.len() + (4 - data.len() % 4), 0);
        for f in self.fns {
            let addr = Addr(data.len() as u64);
            let mut added = vec![(f.sym, addr)];
            sym_table.insert(f.sym, addr);
            for (i, instr) in f.instrs.into_iter().enumerate() {
                let i_pos = Addr(data.len() as u64);
                if let Some(sym) = f.locations.get(&i) {
                    sym_table.insert(*sym, i_pos);
                    added.push((*sym, i_pos));
                }
                if let Some((sym, marker)) = instr.encode(&mut data, &mut sym_table)? {
                    let info = (instr, marker);
                    if let Some(vec) = missing.get_mut(&sym) {
                        vec.push(info);
                    } else {
                        missing.insert(sym, vec![info]);
                    }
                }
            }
            for (sym, addr) in added {
                if let Some(vec) = missing.remove(&sym) {
                    for (i, marker) in vec {
                        i.insert_sym(&mut data, marker, addr);
                    }
                }
            }
        }
        assert!(missing.is_empty());
        Ok(LinkedProgram {
            code: data,
            start: self
                .start
                .map(|s| sym_table.get(s).expect("start symbol doesn't exist")),
        })
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Addr(u64);
impl Addr {
    pub const ZERO: Self = Self(!0);
    pub fn val(&self) -> u64 {
        self.0
    }
}

pub struct SymTable(Vec<Addr>);
impl SymTable {
    pub fn new(len: usize) -> Self {
        Self(vec![Addr::ZERO; len])
    }
    pub fn insert(&mut self, sym: Symbol, addr: Addr) {
        self.0[sym.val()] = addr;
    }
    pub fn get(&self, sym: Symbol) -> Option<Addr> {
        match self.0[sym.val()] {
            Addr::ZERO => None,
            addr => Some(addr),
        }
    }
}

impl LinkedProgram {
    pub fn to_elf(&self) -> Vec<u8> {
        elf::create(&self.code, self.start.expect("no start"))
    }
}