//! This module takes the nested QuickJS binary structure and separates the functions into a flat indexed namespace

use crate::parse::{
    Atom, AtomRef, ClosureKind, Function, Instr, JsString, JsValue, QuickJSBinary, atom_names,
};

#[derive(Debug, Clone, Copy)]
pub enum CaptureFrom {
    /// `parent.var_defs[i]`
    ParentArg(usize),
    /// `parent.var_defs[parent.arg_count + i]`
    ParentLocal(usize),
    /// `parent.captures[i]`, itself resolved against the grandparent.
    ParentCapture(usize),
    /// A module-level binding of the module body.
    ModuleBinding(usize),
    /// Global or module import; not frame-relative.
    NotFrameRelative(&'static str),
}

#[derive(Debug)]
pub struct Capture {
    pub name: String,
    /// The name as it was on the wire. A field cannot be given a name that is
    /// not already an atom, so the operand has to be kept, not rebuilt.
    pub atom: AtomRef,
    pub from: CaptureFrom,
}

#[derive(Debug)]
pub struct FlatFn {
    pub name: String,
    pub arg_count: usize,
    pub var_count: usize,          // num locals
    pub value_names: Vec<AtomRef>, // adds names to args and locals, one for each arg, then one for each local, in one vec
    pub captures: Vec<Capture>,
    pub cpool: Vec<JsValue>,
    pub code: Vec<Instr>,
}

pub fn flatten(bin: QuickJSBinary) -> (Vec<FlatFn>, Vec<Atom>) {
    let mut fns = vec![];
    walk(&mut fns, bin.root_mod.root_fn, &bin.atoms);
    (fns, bin.atoms)
}

fn walk(flat_fns: &mut Vec<FlatFn>, func: Function, atoms: &[Atom]) -> usize {
    let id = flat_fns.len();
    let name = atom_str(&func.name, atoms);

    // Find vars func captures from its environment
    let captures = func
        .closure_vars
        .iter()
        .map(|cv| Capture {
            name: atom_str(&cv.name, atoms),
            atom: cv.name,
            from: match cv.flags.closure_kind {
                ClosureKind::Arg => CaptureFrom::ParentArg(cv.var_idx),
                ClosureKind::Local => CaptureFrom::ParentLocal(cv.var_idx),
                ClosureKind::Ref => CaptureFrom::ParentCapture(cv.var_idx),
                ClosureKind::ModuleDecl => CaptureFrom::ModuleBinding(cv.var_idx),
                ClosureKind::ModuleImport => todo!(), // CaptureFrom::NotFrameRelative("module_import"),
                ClosureKind::Global | ClosureKind::GlobalDecl | ClosureKind::GlobalRef => todo!(), // CaptureFrom::NotFrameRelative("global")
            },
        })
        .collect();

    // Reserve this function's slot before recursing into nested functions,
    // so it keeps `id` while its children take the ids after it
    flat_fns.push(FlatFn {
        name,
        arg_count: func.arg_count,
        var_count: func.var_count,
        value_names: func.var_defs.iter().map(|v| v.name).collect(),
        captures,
        cpool: Vec::new(),
        code: func.byte_code,
    });

    // Nested functions live in the cpool; replace each with its flat id.
    let mut cpool = Vec::with_capacity(func.cpool.len());
    for const_entry in func.cpool {
        let entry = match const_entry {
            JsValue::Function(child) => JsValue::FnById(walk(flat_fns, child, atoms)),
            other => other,
        };
        cpool.push(entry);
    }
    flat_fns[id].cpool = cpool;

    id
}

// LLM GENERATED
// ─── atom rendering ──────────────────────────────────────────────────────

pub fn atom_str(a: &AtomRef, atoms: &[Atom]) -> String {
    match a {
        AtomRef::TaggedInt(n) => n.to_string(),
        AtomRef::Predefined(i) => atom_names::PREDEFINED_ATOMS
            .get(*i as usize)
            .copied()
            .unwrap_or("<predef?>")
            .to_string(),
        AtomRef::FileLocal(i) => atoms
            .get(*i)
            .map(atom_text)
            .unwrap_or_else(|| format!("<file_atom {i}>")),
    }
}

fn atom_text(a: &Atom) -> String {
    match a {
        Atom::String(s) | Atom::GlobalSymbol(s) | Atom::Symbol(s) => js_string(s),
        Atom::Predefined(n) => format!("<predef {n}>"),
    }
}

fn js_string(s: &JsString) -> String {
    match s {
        JsString::Latin1(b) => b.iter().map(|&c| c as char).collect(),
        JsString::UTF16(b) => {
            let u: Vec<u16> = b
                .chunks_exact(2)
                .map(|c| u16::from_le_bytes([c[0], c[1]]))
                .collect();
            String::from_utf16_lossy(&u)
        }
    }
}