use std::collections::HashMap;

use super::{BOX_FIELD, ENV_ARG, FIRST_ARG, pusher};
use crate::{
    libqjs::parse::{
        AtomRef, ClosureKind, ClosureVar, ImmOperand, Instr, JsValue, opcodes::OpCode,
    },
    regify::{is_var_read, is_var_write},
};

/// 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
pub 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,
    }
}