mod declass;

use std::collections::HashMap;

use declass::declass;

use crate::{
    flatten::FlatFn,
    libqjs::parse::{
        Atom, AtomRef, ClosureKind, ClosureVar, ImmOperand, Instr, JsValue, opcodes::OpCode,
    },
    native,
    regify::{is_var_read, is_var_write},
};

/// The field an object's prototype is kept in. In JS, `o.__proto__` is the
/// prototype of an ordinary object, so user code that reads or writes it
/// gets the right behaviour with no extra work.
pub const PROTO_FIELD: AtomRef = AtomRef::Predefined(72); // __proto__

/// The field a constructor keeps the prototype of its instances in
pub const PROTOTYPE_FIELD: AtomRef = AtomRef::Predefined(63); // prototype

/// Every fn takes its env and `this` as its first two params, and its declared
/// params after them. So the params of a call do not depend on how many args
/// the callee declares: a missing arg is not passed, and an extra one is not
/// read.
pub const ENV_ARG: u16 = 0;
pub const THIS_ARG: u16 = 1;
pub const FIRST_ARG: u16 = 2;

/// The field a box keeps its binding's value in (see replace_bindings())
pub const BOX_FIELD: AtomRef = AtomRef::Predefined(86); // <var>

pub fn transform(
    flat_fns: Vec<FlatFn>,
    captures: &[Vec<ClosureVar>],
    atoms: &[Atom],
) -> Vec<FlatFn> {
    flat_fns
        .into_iter()
        .map(|f| transform_one(f, captures, atoms))
        .collect()
}

fn transform_one(mut f: FlatFn, captures: &[Vec<ClosureVar>], atoms: &[Atom]) -> FlatFn {
    let code = declass(&f.code);

    // Temp locals, after the fn's own, for the args of a `new` (see replace_new())
    let temps = f.var_count;
    let n_temps = code.iter().filter_map(is_new).max().unwrap_or(0);

    let mut out = Vec::with_capacity(code.len());

    for instr in &code {
        if replace_new(instr, temps, &mut out) {
            continue;
        }

        if replace_arg_access(instr, &mut out) {
            continue;
        }

        if replace_expansion(instr, &mut out) {
            continue;
        }

        if replace_get_var(instr, atoms, &mut f.cpool, &mut out) {
            continue;
        }

        out.push(instr.clone());
    }

    // Last, so that it also sees the var accesses that the expansions make
    let (code, var_count) = replace_bindings(&out, &f.cpool, captures, temps + n_temps);

    FlatFn {
        name: f.name,
        arg_count: f.arg_count + FIRST_ARG as usize,
        var_count,
        cpool: f.cpool,
        code,
    }
}

/// Pushes instrs at the pc of the instr they replace
fn pusher(pc: usize, out: &mut Vec<Instr>) -> impl FnMut(OpCode, ImmOperand) + '_ {
    move |op, imm| out.push(Instr { pc, op, imm })
}

// ───────────────────────────────────────────────── New Processing ──────────────────────────────────────────────────────────────

/// If `instr` is a `new`, the number of args it passes
fn is_new(instr: &Instr) -> Option<usize> {
    match (instr.op, instr.imm) {
        (OpCode::CallConstructor, ImmOperand::Npop(n)) => Some(n as usize),
        _ => None,
    }
}

/// Replace:
/// call_constructor n                  [F, new_target, a0 .. an-1]
///
/// With:
/// put_loc t(n-1) .. put_loc t0        the args, into temp locals, so F is on top
/// drop                                new_target, always a copy of F here
/// dup                                 [F, F]
/// get_field prototype                 [F, proto]
/// object                              [F, proto, obj]
/// swap                                [F, obj, proto]
/// define_field __proto__              [F, obj], obj is { __proto__: F.prototype }
/// get_loc t0 .. get_loc t(n-1)        the args back
/// call_constructor n                  [F, obj, a0 .. an-1], a method call of F on obj
///
/// Whether `new` gives the call's result or `obj` depends on types (see infer)
fn replace_new(instr: &Instr, temps: usize, out: &mut Vec<Instr>) -> bool {
    let Some(n) = is_new(instr) else {
        return false;
    };

    let mut push = pusher(instr.pc, out);
    let temp = |i: usize| ImmOperand::Loc((temps + i) as u16);

    for i in (0..n).rev() {
        push(OpCode::PutLoc, temp(i));
    }

    push(OpCode::Drop, ImmOperand::None);
    push(OpCode::Dup, ImmOperand::None);
    push(OpCode::GetField, ImmOperand::Atom(PROTOTYPE_FIELD));
    push(OpCode::Object, ImmOperand::None);
    push(OpCode::Swap, ImmOperand::None);
    push(OpCode::DefineField, ImmOperand::Atom(PROTO_FIELD));

    for i in 0..n {
        push(OpCode::GetLoc, temp(i));
    }

    push(OpCode::CallConstructor, instr.imm);
    true
}

// ─────────────────────────────────────────────── Binding Processing ────────────────────────────────────────────────────────────

/// Put every binding that a closure captures in a box: an object whose
/// `BOX_FIELD` holds the binding's value (see `docs/model.md`, section 2). The
/// owner and the closures read and write the box. A closure keeps the boxes
/// it captured in its own fields, one for each of its `closure_vars`: field `i`
/// holds the box of capture `i`
///
/// The fn makes a box for each binding that its children capture, at entry.
/// A captured local holds its box in its own slot. A captured arg keeps its
/// box in a new local, and the box starts with the arg's value. A `var` starts
/// as undefined. A `let` or `const` has no value until it is written (its TDZ):
/// object; undefined; define_field <var>; put_loc k     each captured `var`
/// object; put_loc k                                    each captured `let`
/// object; get_arg a; define_field <var>; put_loc box   each captured arg
///
/// A scope that is entered again is a loop body. QuickJS closes each captured
/// binding of the scope (`close_loc`) at every way out of it, so a new box
/// there gives each iteration its own binding. `set_loc_uninitialized` only
/// marks the TDZ, which is not checked, and hoisted fns are made before it, so
/// it does not make a box
///
/// These instrs take a pc that no jump goes to, so a jump to the fn's first
/// instr does not make them again
///
/// Then it replaces, for a captured binding:
/// get_loc k                   get_loc box; get_field <var>
/// put_loc k                   get_loc box; swap; put_field <var>
/// set_loc k                   dup; get_loc box; swap; put_field <var>
/// close_loc k                 object; get_loc k; get_field <var>;
///                             define_field <var>; put_loc k
///                             a way out of a loop body: a new binding, with
///                             the current value
///
/// For any binding:
/// get_var_ref i               get_arg env; get_field i; get_field <var>
/// put_var_ref i               get_arg env; get_field i; swap; put_field <var>
/// set_var_ref i               dup; get_arg env; get_field i; swap; put_field <var>
/// fclosure f                  fclosure f; (push_box; define_field i)*
///
/// where `push_box` is the box of capture `i` of `f`: the local that holds it,
/// or field `j` of this fn's env for capture `j` of this fn. A binding with no
/// frame (a module binding) gets its box here, from `object`
///
/// `set_loc_uninitialized`, and a `close_loc` of a binding no closure
/// captures, do nothing. Gives the code and the new number of locals
fn replace_bindings(
    code: &[Instr],
    cpool: &[JsValue],
    captures: &[Vec<ClosureVar>],
    var_count: usize,
) -> (Vec<Instr>, usize) {
    use OpCode::*;

    let children = cpool.iter().filter_map(|c| match c {
        JsValue::FnById(child) => Some(*child),
        _ => None,
    });

    let mut locals = vec![];
    let mut args = vec![];

    for c in children.flat_map(|child| &captures[child]) {
        match c.flags.closure_kind {
            ClosureKind::Local => locals.push((c.var_idx, c.flags.is_lexical)),
            ClosureKind::Arg => args.push(c.var_idx),
            _ => {}
        }
    }

    locals.sort_unstable();
    locals.dedup();
    args.sort_unstable();
    args.dedup();

    // Value ids of bindings, as `is_var_read` gives them: args after every local
    let new_var_count = var_count + args.len();
    let arg_id = |k: usize| new_var_count + FIRST_ARG as usize + k;

    // The local that holds the box of each captured binding, by value id
    let boxes: HashMap<usize, usize> = (locals.iter().map(|&(k, _)| (k, k)))
        .chain(args.iter().enumerate().map(|(i, &k)| (arg_id(k), var_count + i)))
        .collect();

    let slot = |i: usize| ImmOperand::Atom(AtomRef::TaggedInt(i as u32));
    let value = ImmOperand::Atom(BOX_FIELD);
    let none = ImmOperand::None;
    let loc = |k: usize| ImmOperand::Loc(k as u16);

    let mut out = Vec::with_capacity(code.len());

    {
        let mut push = pusher(usize::MAX, &mut out);

        for &(k, lexical) in &locals {
            push(Object, none);

            if !lexical {
                push(Undefined, none);
                push(DefineField, value);
            }

            push(PutLoc, loc(k));
        }

        for (i, &k) in args.iter().enumerate() {
            push(Object, none);
            push(GetArg, ImmOperand::Arg(FIRST_ARG + k as u16));
            push(DefineField, value);
            push(PutLoc, loc(var_count + i));
        }
    }

    for instr in code {
        let read = is_var_read(instr, new_var_count).and_then(|v| boxes.get(&v));
        let write = is_var_write(instr, new_var_count).and_then(|v| boxes.get(&v));
        let own_box = |k: ImmOperand| match k {
            ImmOperand::Loc(k) => boxes.get(&(k as usize)).copied(),
            _ => None,
        };
        let mut push = pusher(instr.pc, &mut out);

        if let Some(&b) = read {
            push(GetLoc, loc(b));
            push(GetField, value);
        } else if let Some(&b) = write {
            if instr.op.info().n_push == 1 {
                push(Dup, none);
            }

            push(GetLoc, loc(b));
            push(Swap, none);
            push(PutField, value);
        } else if let Some((i, writes, leaves)) = var_ref_access(instr) {
            if leaves {
                push(Dup, none);
            }

            push(GetArg, ImmOperand::Arg(ENV_ARG));
            push(GetField, slot(i));

            if writes {
                push(Swap, none);
                push(PutField, value);
            } else {
                push(GetField, value);
            }
        } else if instr.op == SetLocUninitialized {
            push(Nop, none);
        } else if instr.op == CloseLoc {
            match own_box(instr.imm) {
                Some(b) => {
                    push(Object, none);
                    push(GetLoc, loc(b));
                    push(GetField, value);
                    push(DefineField, value);
                    push(PutLoc, loc(b));
                }
                None => push(Nop, none),
            }
        } else if matches!(instr.op, Fclosure8 | Fclosure) {
            let JsValue::FnById(child) = cpool[instr.imm.const_idx()] else {
                unreachable!()
            };

            push(Fclosure, instr.imm);

            for (i, c) in captures[child].iter().enumerate() {
                match c.flags.closure_kind {
                    ClosureKind::Local => push(GetLoc, loc(boxes[&c.var_idx])),
                    ClosureKind::Arg => push(GetLoc, loc(boxes[&arg_id(c.var_idx)])),
                    ClosureKind::Ref => {
                        push(GetArg, ImmOperand::Arg(ENV_ARG));
                        push(GetField, slot(c.var_idx));
                    }
                    _ => push(Object, none),
                }

                push(DefineField, slot(i));
            }
        } else {
            push(instr.op, instr.imm);
        }
    }

    (out, new_var_count)
}

/// If `instr` reads or writes capture `i` of its fn: `i`, whether it writes,
/// and whether a write also leaves the value on the stack
fn var_ref_access(instr: &Instr) -> Option<(usize, bool, bool)> {
    use OpCode::*;

    let i = match instr.imm {
        ImmOperand::VarRef(i) => i as usize,
        _ => 0,
    };

    match instr.op {
        GetVarRef0 | GetVarRef1 | GetVarRef2 | GetVarRef3 => Some((instr.op as usize - GetVarRef0 as usize, false, false)),
        PutVarRef0 | PutVarRef1 | PutVarRef2 | PutVarRef3 => Some((instr.op as usize - PutVarRef0 as usize, true, false)),
        SetVarRef0 | SetVarRef1 | SetVarRef2 | SetVarRef3 => Some((instr.op as usize - SetVarRef0 as usize, true, true)),
        GetVarRef | GetVarRefCheck => Some((i, false, false)),
        PutVarRef | PutVarRefCheck | PutVarRefCheckInit => Some((i, true, false)),
        SetVarRef => Some((i, true, true)),
        _ => None,
    }
}

// ──────────────────────────────────────────────── Arg Access Processing ────────────────────────────────────────────────────────

/// Replace an access of declared param `k` with one of param `FIRST_ARG + k`,
/// and `this` with a read of its param
fn replace_arg_access(instr: &Instr, out: &mut Vec<Instr>) -> bool {
    use OpCode::*;

    let (op, arg) = match (instr.op, instr.imm) {
        (PushThis, _) => (GetArg, THIS_ARG),
        (GetArg0, _) => (GetArg, FIRST_ARG),
        (GetArg1, _) => (GetArg, FIRST_ARG + 1),
        (GetArg2, _) => (GetArg, FIRST_ARG + 2),
        (GetArg3, _) => (GetArg, FIRST_ARG + 3),
        (PutArg0, _) => (PutArg, FIRST_ARG),
        (PutArg1, _) => (PutArg, FIRST_ARG + 1),
        (PutArg2, _) => (PutArg, FIRST_ARG + 2),
        (PutArg3, _) => (PutArg, FIRST_ARG + 3),
        (SetArg0, _) => (SetArg, FIRST_ARG),
        (SetArg1, _) => (SetArg, FIRST_ARG + 1),
        (SetArg2, _) => (SetArg, FIRST_ARG + 2),
        (SetArg3, _) => (SetArg, FIRST_ARG + 3),
        (GetArg | PutArg | SetArg, ImmOperand::Arg(k)) => (instr.op, FIRST_ARG + k),
        _ => return false,
    };

    pusher(instr.pc, out)(op, ImmOperand::Arg(arg));
    true
}

// ───────────────────────────────────────────────── Misc Expansion ──────────────────────────────────────────────────────────────

/// Replace ops that are a fixed sequence of simpler ones
///
/// A sequence is never empty: jumps find their target by its pc (see regify),
/// so every instr must keep at least one instr at its pc
fn replace_expansion(instr: &Instr, out: &mut Vec<Instr>) -> bool {
    use OpCode::*;

    let none = ImmOperand::None;
    let imm = instr.imm;

    // A plain call has no `this`, so it passes `undefined`. It is pushed last,
    // on top of the args, and the call pops it as one more arg (see regify)
    let plain_call = |n: u16| [(Undefined, none), (Call, ImmOperand::Npop(n + 1))];
    let tail_call = |n: u16| [(Undefined, none), (Call, ImmOperand::Npop(n + 1)), (Return, none)];

    let seq: &[(OpCode, ImmOperand)] = match (instr.op, imm) {
        (Call0, _) => &plain_call(0),
        (Call1, _) => &plain_call(1),
        (Call2, _) => &plain_call(2),
        (Call3, _) => &plain_call(3),
        (Call, ImmOperand::Npop(n)) => &plain_call(n),
        (TailCall, ImmOperand::Npop(n)) => &tail_call(n),
        (TailCallMethod, _) => &[(CallMethod, imm), (Return, none)],
        (GetField2, _) => &[(Dup, none), (GetField, imm)],
        (GetLoc0Loc1, _) => &[(GetLoc0, none), (GetLoc1, none)],
        (IncLoc, _) => &[(GetLoc, imm), (Inc, none), (PutLoc, imm)],
        (DecLoc, _) => &[(GetLoc, imm), (Dec, none), (PutLoc, imm)],
        // The added value is on the stack, the local is the left operand
        (AddLoc, _) => &[(GetLoc, imm), (Swap, none), (Add, none), (PutLoc, imm)],
        (SetProto, _) => &[(DefineField, ImmOperand::Atom(PROTO_FIELD))],
        // The low two bits of `arg`: 0 method, 1 getter, 2 setter
        (DefineMethod, ImmOperand::AtomU8 { atom, arg }) => {
            assert!(arg & 3 == 0, "getters and setters are not supported");
            &[(DefineField, ImmOperand::Atom(atom))]
        }
        // The fn itself, which is its env
        (SpecialObject, ImmOperand::U8(2)) => &[(GetArg, ImmOperand::Arg(ENV_ARG))],
        // [[HomeObject]], which only `super` reads. `super` is not supported
        // (see set_home_object in regify)
        (SpecialObject, ImmOperand::U8(4)) => &[(Undefined, none)],
        // arguments, new.target, import.meta
        (SpecialObject, _) => unimplemented!("special_object {imm:?}"),
        _ => return false,
    };

    let mut push = pusher(instr.pc, out);
    seq.iter().for_each(|&(op, imm)| push(op, imm));
    true
}

/// Replace a read of a global with the constant it is (see `native`)
fn replace_get_var(
    instr: &Instr,
    atoms: &[Atom],
    cpool: &mut Vec<JsValue>,
    out: &mut Vec<Instr>,
) -> bool {
    if instr.op != OpCode::GetVar {
        return false;
    }

    cpool.push(native::global(instr.imm.atom(), atoms));
    pusher(instr.pc, out)(OpCode::PushConst, ImmOperand::Const(cpool.len() as u32 - 1));
    true
}