use std::collections::{HashMap, HashSet, VecDeque};

use super::*;
use crate::backend::{Body, Func, LibImport, LinkedProgram, SymImport, SymTable, Symbol, VarId};

pub struct Encoder<'a> {
    pub code: Code,
    pub sym_tab: SymTable<u64>,
    pub sym_refs: HashMap<Symbol, Vec<usize>>,
    pub program: &'a Program<X86_64>,
    pub segs: Vec<SegUses>,
    pub vars: HashMap<VarId, Reg>,
}

pub fn compile(p: &Program<X86_64>) -> Result<LinkedProgram<u64>, CompilerMsg> {
    let mut encoder = Encoder::new(p);

    p.encode_data(&mut encoder.code.bytes, &mut encoder.sym_tab);

    for f in &p.funcs {
        encoder.func(f);
    }

    for (pos, sym) in encoder.code.missing.drain(..) {
        let info = encoder.program.sym_info(sym);
        if info.external {
            encoder.sym_refs.entry(sym).or_default().push(pos);
        } else {
            let addr = encoder
                .sym_tab
                .get(sym)
                .ok_or(CompilerMsg::from(format!("missing symbol {}", info.name)))?;
            encoder.code.bytes[pos..pos + 4].copy_from_slice(&addr_offset(pos, addr))
        }
    }

    let imports = p
        .external
        .iter()
        .map(|e| LibImport {
            name: e.file.clone(),
            syms: e
                .syms
                .iter()
                .map(|&s| SymImport {
                    name: p.sym_info(s).name.clone(),
                    usages: encoder.sym_refs.entry(s).or_default().clone(),
                })
                .collect(),
        })
        .collect();

    Ok(LinkedProgram {
        code: encoder.code.bytes,
        entry: p.entry.and_then(|e| encoder.sym_tab.get(e)),
        imports,
    })
}

#[derive(Default)]
pub struct SegUses {
    var: HashMap<VarId, VecDeque<usize>>,
    reg: [VecDeque<RegUse>; 16],
}

#[derive(Clone, Copy)]
pub struct RegUse {
    pos: usize,
    // required var, if any
    var: Option<VarId>,
}

type BInstr = crate::backend::Instr<X86_64>;
impl<'a> Encoder<'a> {
    fn func(&mut self, f: &Func<X86_64>) -> EncodeRes {
        let addr = self.code.bytes.len();
        self.sym_tab.insert(f.sym, addr as u64);
        self.segs.clear();
        self.calc_uses(&f.body);
        for instr in &f.body {
            match instr {
                BInstr::Asm(asm) => self.asm(asm)?,
                _ => todo!(),
            }
        }
        Ok(())
    }

    pub fn calc_uses(&mut self, body: &Body<X86_64>) -> (usize, usize) {
        let mut pos = 0;
        let seg_i = self.segs.len();
        self.segs.push(Default::default());
        let mut uses = SegUses::default();
        macro_rules! push {
            ($var:ident) => {
                push!($var, 0)
            };
            ($var:ident, $pos:expr) => {
                uses.var
                    .entry($var.clone())
                    .or_default()
                    .push_back(pos + $pos)
            };
        }
        for instr in body {
            match instr {
                BInstr::Set { dst, src: _ } => push!(dst),
                BInstr::Call { dst, f, args } => {
                    let conv = &self.program.call_convs[self.program.funcs[f].conv];
                    push!(dst);

                    for &reg in conv.scratch() {
                        uses.reg[reg as usize].push_back(RegUse { pos, var: None });
                    }
                    for (i, &arg) in args.iter().enumerate() {
                        push!(arg);
                        if let Some(&reg) = conv.param().get(i) {
                            uses.reg[reg as usize].back_mut().unwrap().var = Some(arg);
                        }
                    }
                }
                BInstr::Copy { dst, src } => {
                    push!(dst);
                    push!(src);
                }
                BInstr::Add { dst, src1, src2 } => {
                    push!(dst);
                    push!(src1);
                    push!(src2);
                }
                BInstr::If { cond, then, else_ } => {
                    push!(cond);
                    pos += 1;
                    let (seg_i1, len1) = self.calc_uses(then);
                    let (seg_i2, len2) = self.calc_uses(else_);
                    // insert closest usages
                    for (var, poss1) in &self.segs[seg_i1].var {
                        if let Some(poss2) = self.segs[seg_i2].var.get(var) {
                            push!(var, poss1[0].min(poss2[0]));
                        } else {
                            push!(var, poss1[0]);
                        }
                    }
                    for (var, poss2) in &self.segs[seg_i2].var {
                        if !self.segs[seg_i1].var.contains_key(var) {
                            push!(var, poss2[0]);
                        }
                    }
                    pos += len1.max(len2);
                    continue;
                }
                BInstr::Loop(instrs) => {
                    let (seg_i2, len) = self.calc_uses(instrs);
                    // insert closest usages
                    for (var, poss) in &self.segs[seg_i2].var {
                        push!(var, poss[0]);
                    }
                    // during register allocation, want to insert the first use of each var
                    // at len + use, because if it jumps back up, that will be the next
                    // usage rather than if the loop exits; pad ensures there is room so
                    // the insertions are always before the rest of the usages (after loop)
                    let pad = (len * 2).saturating_sub(pos);
                    pos += len + pad;
                    continue;
                }
                BInstr::Break(var) => {
                    if let Some(var) = var {
                        push!(var);
                    }
                }
                BInstr::Return(var) => {
                    if let Some(var) = var {
                        push!(var);
                    }
                }
                BInstr::Asm(asm) => {
                    let mut used = HashSet::new();
                    for i in &asm.instrs {
                        if let Some(reg) = i.dst_reg() {
                            used.insert(reg.reg());
                        }
                    }
                    for &(reg, var) in &asm.args {
                        push!(var);
                        used.remove(&reg);
                        uses.reg[reg as usize].push_back(RegUse {
                            pos,
                            var: Some(var),
                        });
                    }
                    for reg in used {
                        uses.reg[reg as usize].push_back(RegUse { pos, var: None });
                    }
                }
            }
            pos += 1;
        }
        self.segs[seg_i] = uses;
        return (seg_i, pos);
    }

    pub fn rvm(&mut self, input: Rvm) -> RegMemKind {
        match input {
            Rvm::Reg(reg) => reg,
            Rvm::Var(var_id) => todo!(),
            Rvm::Mem(mem) => todo!(),
        }
    }

    pub fn asm(&mut self, asm: &Asm) -> EncodeRes {
        let mut used = HashSet::default();
        let mut vars = HashSet::default();
        for &instr in &asm.instrs {
            if let Some(regw) = instr.dst_reg() {
                used.insert(regw.reg());
            }
            for var in instr.vars() {
                vars.insert(var);
            }
        }
        let overlap = used.intersection(&self.vars.values());
        for var in &vars {
            if let Some(reg) = self.vars.get(var)
                && used.contains(reg)
            {}
        }
        for &instr in &asm.instrs {
            match instr {
                Instr::Mov { dst, src } => self.code.mov(dst, src)?,
                Instr::Push(rvmi) => todo!(),
                Instr::Pop(rvm) => todo!(),
                Instr::Lea { dst, src } => todo!(),
                Instr::Int(code) => self.code.int(code),
                Instr::Syscall => self.code.syscall(),
                Instr::Call(sym) => self.code.call(sym),
                Instr::CallAt(sym) => self.code.call_mem(sym),
                Instr::Ret => self.code.ret(),
                Instr::Add { dst, src } => todo!(),
                Instr::Sub { dst, src } => todo!(),
            }
        }
        Ok(())
    }

    pub fn new(program: &'a Program<X86_64>) -> Self {
        Self {
            code: Code::default(),
            sym_tab: SymTable::new(program.sym_count()),
            sym_refs: Default::default(),
            segs: Default::default(),
            program,
        }
    }
}