use std::range::Range;

use crate::{
    parse::{AtomRef, ImmOperand, JsValue, atom_names::PREDEFINED_ATOMS, opcodes::OpCode},
    regify::{RegFn, RegInstr},
};

pub fn find_class_seqs(reg_fn: &RegFn) -> Vec<ClassSequence> {
    let mut ins = reg_fn.code.as_slice();
    let mut seqs = vec![];
    let mut offset = 0;

    while !ins.is_empty() {
        let seq = match_class_seq(ins); // Try to match a class decl sequence

        match seq {
            // Found class sequence, store it and advance `ins`
            Ok((rest, num_methods)) => {
                let consumed = ins.len() - rest.len(); // How many instrs were consumed
                let seq_instr = &ins[..consumed]; // The class seq's instrs

                seqs.push(ClassSequence {
                    range: Range {
                        start: offset,
                        end: offset + consumed,
                    },
                    num_methods,
                    ctor_fn: fn_id(&reg_fn.cpool, as_const(seq_instr[0].imm)),
                    field_init_fn: find_field_init_fn(seq_instr, num_methods, &reg_fn.cpool),
                });

                offset += consumed;
                ins = rest;
            }
            // Did not find sequence, advance `ins` by 1 and try again
            Err(rest) => {
                let consumed = ins.len() - rest.len(); // Handle partial consumption
                let consumed = consumed.max(1);
                offset += consumed;
                ins = &ins[consumed..];
            }
        }
    }

    seqs
}

pub fn rewrite_class_to_proto(reg_fn: &mut RegFn, class_seqs: &[ClassSequence]) {
    // The prototype form is never longer than the class build it replaces, so
    // it drops straight over it and every instr outside a class seq stays put.
    // The code moves out only so fresh_val() can borrow reg_fn beside it.
    let mut code = std::mem::take(&mut reg_fn.code);
    let mut buf = vec![]; // Reused across seqs, so it allocates once

    for seq in class_seqs {
        seq.emit_prototype(&code[seq.range], &mut || reg_fn.fresh_val(), &mut buf);
        code[seq.range].swap_with_slice(&mut buf);
    }

    reg_fn.code = code;
}

pub struct ClassSequence {
    pub range: Range<usize>,
    pub num_methods: usize,
    pub ctor_fn: usize,               // Fn id of the class ctor
    pub field_init_fn: Option<usize>, // Fn id of the field init, if any
}

/// The field initialiser's fn id, if the class has a field-initialiser tail.
fn find_field_init_fn(
    matched: &[RegInstr],
    num_methods: usize,
    cpool: &[JsValue],
) -> Option<usize> {
    let tail = 2 + 2 * num_methods; // The ctor and its methods, before the tail

    if matched.len() == tail {
        return None; // No field-initialiser tail
    }

    Some(fn_id(cpool, as_const(matched[tail].imm)))
}

impl ClassSequence {
    /// Emit the prototype form into the scratchpad `buf`, reading operands
    /// straight out of `matched` -- this class sequence's own instructions.
    /// `buf` comes back padded with nops to exactly the length of the class
    /// build, so the caller can drop it over the original in place.
    fn emit_prototype<F>(&self, matched: &[RegInstr], fresh_val: &mut F, buf: &mut Vec<RegInstr>)
    where
        F: FnMut() -> usize,
    {
        buf.clear();

        // push_const8 #ctor ; define_class  ->  fclosure8 #ctor (constructor kept)
        let (ctor_const, ctor_val) = self.ctor(matched);
        buf.push(fclosure8(ctor_const, ctor_val));

        // A class with no methods has no prototype to read
        if self.num_methods > 0 {
            // The prototype cannot change between the put_fields below, so read
            // it once for the whole class instead of once per method
            let proto = fresh_val();
            buf.push(get_field(proto_atom(), proto, ctor_val));

            // each ( fclosure8 #m ; define_method @meth ) -> fclosure8 + put_field
            for k in 0..self.num_methods {
                let (method_const, method_val, meth) = self.method(matched, k);
                buf.push(fclosure8(method_const, method_val));
                buf.push(put_field(ImmOperand::Atom(meth), method_val, proto));
            }
        }
        // the fclosure8 #fields ; set_home_object ; put_loc ; drop tail is dropped.

        // 2 + 2n emitted against the 2 + 2n matched, less the 4-instr field
        // initialiser tail, so the emission always fits. Nop pads the slack; it
        // pops and pushes nothing, so it cannot disturb the value flow.
        let len = self.range.end - self.range.start;
        debug_assert!(buf.len() <= len, "class rewrite outgrew its range");
        buf.resize_with(len, nop);
    }

    /// The constructor's constant and register ID it was written to
    /// push_const8 #ctor ; define_class.
    fn ctor(&self, matched: &[RegInstr]) -> (u32, usize) {
        (as_const(matched[0].imm), matched[1].dsts[0])
    }

    /// The k-th method's constant, closure value, and name; each method pair
    /// sits after the two constructor instructions.
    fn method(&self, matched: &[RegInstr], k: usize) -> (u32, usize, AtomRef) {
        let fclosure = &matched[2 + 2 * k];
        let define_method = &matched[3 + 2 * k];
        (
            as_const(fclosure.imm),
            fclosure.dsts[0],
            atom_of(define_method.imm),
        )
    }
}

// ─────────────────────────── Class Sequence Grammar ─────────────────────────────────────
// push_const8 #ctor ; define_class @Name                       <-- Class decl header
// ( fclosure8 #method ; define_method @meth )*                 <-- Methods
// ( fclosure8 #fields ; set_home_object ; put_loc ; drop )?    <-- Field initializer stuff
//
// match fns return Ok(advanced_instrs) or Err(original_instrs)

fn match_class_seq(ins: &[RegInstr]) -> Result<(&[RegInstr], usize), &[RegInstr]> {
    let ins = match_one(ins, |i| i.op == OpCode::PushConst8)?;

    if let Some(i) = ins.first() {
        let is_derived_class = matches!(i.imm, ImmOperand::AtomU8 { arg, .. } if arg & 1 == 1);

        if i.op == OpCode::DefineClass && is_derived_class {
            panic!("declass: derived classes (extends) are not supported");
        }
    }

    let ins = match_one(ins, |i| i.op == OpCode::DefineClass)?;
    let (ins, num_methods) = match_many(ins, match_method_seq);
    let ins = match_opt(ins, match_field_init_seq);

    Ok((ins, num_methods))
}

// TODO: accessors are not handled
fn match_method_seq(ins: &[RegInstr]) -> Result<&[RegInstr], &[RegInstr]> {
    let ins = match_one(ins, |i| i.op == OpCode::Fclosure8)?;
    let ins = match_one(ins, |i| i.op == OpCode::DefineMethod)?;
    Ok(ins)
}

fn match_field_init_seq(ins: &[RegInstr]) -> Result<&[RegInstr], &[RegInstr]> {
    let is_put_loc = |op| {
        matches!(
            op,
            OpCode::PutLoc
                | OpCode::PutLoc8
                | OpCode::PutLoc0
                | OpCode::PutLoc1
                | OpCode::PutLoc2
                | OpCode::PutLoc3
        )
    };

    let ins = match_one(ins, |i| i.op == OpCode::Fclosure8)?;
    let ins = match_one(ins, |i| i.op == OpCode::SetHomeObject)?;
    let ins = match_one(ins, |i| is_put_loc(i.op))?;
    let ins = match_one(ins, |i| i.op == OpCode::Drop)?;
    Ok(ins)
}

// ───────────────────────────── Matcher Primitives ───────────────────────────────────────

fn match_one(i: &[RegInstr], f: impl Fn(&RegInstr) -> bool) -> Result<&[RegInstr], &[RegInstr]> {
    match i.first() {
        Some(instr) if f(instr) => Ok(&i[1..]),
        _ => Err(i),
    }
}

fn match_many(
    mut i: &[RegInstr],
    f: impl Fn(&[RegInstr]) -> Result<&[RegInstr], &[RegInstr]>,
) -> (&[RegInstr], usize) {
    let mut n = 0;

    while let Ok(rest) = f(i) {
        n += 1;
        i = rest;
    }

    (i, n)
}

fn match_opt<'a>(
    i: &'a [RegInstr],
    f: impl Fn(&'a [RegInstr]) -> Result<&'a [RegInstr], &'a [RegInstr]>,
) -> &'a [RegInstr] {
    match f(i) {
        Ok(rest) => rest,
        Err(_) => i,
    }
}

// ── emitted instructions ───────────────────────────────────────────────────

fn nop() -> RegInstr {
    RegInstr::new(OpCode::Nop, ImmOperand::None, vec![], vec![])
}

fn fclosure8(const_index: u32, dst: usize) -> RegInstr {
    RegInstr::new(
        OpCode::Fclosure8,
        ImmOperand::Const(const_index),
        vec![dst],
        vec![],
    )
}

fn get_field(field: ImmOperand, dst: usize, obj: usize) -> RegInstr {
    RegInstr::new(OpCode::GetField, field, vec![dst], vec![obj])
}

fn put_field(field: ImmOperand, value: usize, obj: usize) -> RegInstr {
    RegInstr::new(OpCode::PutField, field, vec![], vec![value, obj])
}

// ── operand helpers ────────────────────────────────────────────────────────

fn atom_of(imm: ImmOperand) -> AtomRef {
    match imm {
        ImmOperand::AtomU8 { atom, .. } => atom,
        _ => unreachable!("expected an atom operand"),
    }
}

fn as_const(imm: ImmOperand) -> u32 {
    match imm {
        ImmOperand::Const(c) => c,
        _ => unreachable!("expected a constant operand"),
    }
}

/// The function id a constant-pool entry points at.
fn fn_id(cpool: &[JsValue], const_index: u32) -> usize {
    match cpool[const_index as usize] {
        JsValue::FnById(id) => id,
        _ => unreachable!("class constant is not a function"),
    }
}

/// The `"prototype"` field name as a predefined-atom operand.
fn proto_atom() -> ImmOperand {
    let i = PREDEFINED_ATOMS
        .iter()
        .position(|&n| n == "prototype")
        .unwrap();
    ImmOperand::Atom(AtomRef::Predefined(i as u32))
}