use fix_bytecode::PrimOpPhase;
use fix_runtime::{
    AttrSet, BytecodeReader, CallFrame, List, Machine, MachineExt, NixNum, Null, Path, Step,
    StrictValue, Value, VmRuntimeCtx, VmRuntimeCtxExt,
};
use gc_arena::{Gc, Mutation};
use smallvec::SmallVec;

pub fn start_eq<'gc, M: Machine<'gc>>(
    m: &mut M,
    ctx: &impl VmRuntimeCtx,
    reader: &mut BytecodeReader<'_>,
    mc: &Mutation<'gc>,
    lhs: StrictValue<'gc>,
    rhs: StrictValue<'gc>,
    negate: bool,
) -> Step {
    match shallow_eq(ctx, lhs, rhs) {
        ShallowEq::True => {
            m.push(Value::new_inline(!negate));
            Step::Continue(())
        }
        ShallowEq::False => {
            m.push(Value::new_inline(negate));
            Step::Continue(())
        }
        ShallowEq::RecurseList(la, lb) => {
            let lhs_init: SmallVec<[Value<'gc>; 4]> = la.inner.borrow().iter().copied().collect();
            let rhs_init: SmallVec<[Value<'gc>; 4]> = lb.inner.borrow().iter().copied().collect();
            enter_eq_machine(m, reader, mc, negate, lhs_init, rhs_init)
        }
        ShallowEq::RecurseAttrs(a, b) => {
            let lhs_init: SmallVec<[Value<'gc>; 4]> = a.entries.iter().map(|&(_, v)| v).collect();
            let rhs_init: SmallVec<[Value<'gc>; 4]> = b.entries.iter().map(|&(_, v)| v).collect();
            enter_eq_machine(m, reader, mc, negate, lhs_init, rhs_init)
        }
    }
}

pub fn eq_step<'gc, M: Machine<'gc>>(
    m: &mut M,
    reader: &mut BytecodeReader<'_>,
    mc: &Mutation<'gc>,
) -> Step {
    let rhs_q = m
        .peek(0)
        .as_gc::<List<'gc>>()
        .expect("eq state corrupted: rhs_queue");
    let lhs_q = m
        .peek(1)
        .as_gc::<List<'gc>>()
        .expect("eq state corrupted: lhs_queue");
    let result = m
        .peek(2)
        .as_inline::<bool>()
        .expect("eq state corrupted: result");

    if !result || lhs_q.inner.borrow().is_empty() {
        return finalize(m, reader);
    }

    let lhs = lhs_q
        .unlock(mc)
        .borrow_mut()
        .pop()
        .expect("non-empty lhs queue");
    let rhs = rhs_q
        .unlock(mc)
        .borrow_mut()
        .pop()
        .expect("non-empty rhs queue");
    m.push(lhs);
    m.push(rhs);
    reader.set_pc(PrimOpPhase::EqForce.ip() as usize);
    Step::Continue(())
}

pub fn eq_force<'gc, M: Machine<'gc>>(
    m: &mut M,
    ctx: &mut impl VmRuntimeCtx,
    reader: &mut BytecodeReader<'_>,
    mc: &Mutation<'gc>,
) -> Step {
    let (lhs, rhs) = m.force_and_retry::<(StrictValue, StrictValue)>(reader, mc)?;
    apply_pair(m, ctx, mc, lhs, rhs);
    reader.set_pc(PrimOpPhase::EqStep.ip() as usize);
    Step::Continue(())
}

fn finalize<'gc, M: Machine<'gc>>(m: &mut M, reader: &mut BytecodeReader<'_>) -> Step {
    let _ = m.pop();
    let _ = m.pop();
    let result = m
        .pop()
        .as_inline::<bool>()
        .expect("eq state corrupted: result");
    let negate = m
        .pop()
        .as_inline::<bool>()
        .expect("eq state corrupted: negate");
    m.return_from_primop(Value::new_inline(result ^ negate), reader)
}

fn apply_pair<'gc, M: Machine<'gc>>(
    m: &mut M,
    ctx: &impl VmRuntimeCtx,
    mc: &Mutation<'gc>,
    lhs: StrictValue<'gc>,
    rhs: StrictValue<'gc>,
) {
    match shallow_eq(ctx, lhs, rhs) {
        ShallowEq::True => {}
        ShallowEq::False => {
            m.replace(2, Value::new_inline(false));
        }
        ShallowEq::RecurseList(la, lb) => {
            extend_queues(
                m,
                mc,
                la.inner.borrow().iter().copied(),
                lb.inner.borrow().iter().copied(),
            );
        }
        ShallowEq::RecurseAttrs(a, b) => {
            extend_queues(
                m,
                mc,
                a.entries.iter().map(|&(_, v)| v),
                b.entries.iter().map(|&(_, v)| v),
            );
        }
    }
}

fn extend_queues<'gc, M, L, R>(m: &mut M, mc: &Mutation<'gc>, lhs_iter: L, rhs_iter: R)
where
    M: Machine<'gc>,
    L: IntoIterator<Item = Value<'gc>>,
    R: IntoIterator<Item = Value<'gc>>,
{
    let rhs_q = m
        .peek(0)
        .as_gc::<List<'gc>>()
        .expect("eq state corrupted: rhs_queue");
    let lhs_q = m
        .peek(1)
        .as_gc::<List<'gc>>()
        .expect("eq state corrupted: lhs_queue");
    let mut lq = lhs_q.unlock(mc).borrow_mut();
    let mut rq = rhs_q.unlock(mc).borrow_mut();
    for (x, y) in lhs_iter.into_iter().zip(rhs_iter) {
        lq.push(x);
        rq.push(y);
    }
}

fn enter_eq_machine<'gc, M: Machine<'gc>>(
    m: &mut M,
    reader: &mut BytecodeReader<'_>,
    mc: &Mutation<'gc>,
    negate: bool,
    lhs_init: SmallVec<[Value<'gc>; 4]>,
    rhs_init: SmallVec<[Value<'gc>; 4]>,
) -> Step {
    let resume_pc = reader.pc();
    m.push_call_frame(CallFrame {
        pc: resume_pc,
        thunk: None,
        env: m.env(),
    });
    m.inc_call_depth();
    m.push(Value::new_inline(negate));
    m.push(Value::new_inline(true));
    m.push(Value::new_gc(List::new(mc, lhs_init)));
    m.push(Value::new_gc(List::new(mc, rhs_init)));
    reader.set_pc(PrimOpPhase::EqStep.ip() as usize);
    Step::Continue(())
}

enum ShallowEq<'gc> {
    True,
    False,
    RecurseList(Gc<'gc, List<'gc>>, Gc<'gc, List<'gc>>),
    RecurseAttrs(Gc<'gc, AttrSet<'gc>>, Gc<'gc, AttrSet<'gc>>),
}

fn shallow_eq<'gc>(
    ctx: &impl VmRuntimeCtx,
    lhs: StrictValue<'gc>,
    rhs: StrictValue<'gc>,
) -> ShallowEq<'gc> {
    if let (Some(a), Some(b)) = (lhs.as_num(), rhs.as_num()) {
        let eq = match (a, b) {
            (NixNum::Int(a), NixNum::Int(b)) => a == b,
            (NixNum::Float(a), NixNum::Float(b)) => a == b,
            (NixNum::Int(a), NixNum::Float(b)) => a as f64 == b,
            (NixNum::Float(a), NixNum::Int(b)) => a == b as f64,
        };
        return bool_outcome(eq);
    }
    if let (Some(a), Some(b)) = (lhs.as_inline::<bool>(), rhs.as_inline::<bool>()) {
        return bool_outcome(a == b);
    }
    if lhs.is::<Null>() && rhs.is::<Null>() {
        return ShallowEq::True;
    }
    if let (Some(a), Some(b)) = (lhs.as_inline::<Path>(), rhs.as_inline::<Path>()) {
        return bool_outcome(a.0 == b.0);
    }
    if let (Some(a), Some(b)) = (ctx.get_string(lhs), ctx.get_string(rhs)) {
        return bool_outcome(a == b);
    }
    if let (Some(a), Some(b)) = (lhs.as_gc::<List<'gc>>(), rhs.as_gc::<List<'gc>>()) {
        if a.inner.borrow().len() != b.inner.borrow().len() {
            return ShallowEq::False;
        }
        return ShallowEq::RecurseList(a, b);
    }
    if let (Some(a), Some(b)) = (lhs.as_gc::<AttrSet<'gc>>(), rhs.as_gc::<AttrSet<'gc>>()) {
        let ae = &a.entries;
        let be = &b.entries;
        if ae.len() != be.len() {
            return ShallowEq::False;
        }
        for (l, r) in ae.iter().zip(be.iter()) {
            if l.0 != r.0 {
                return ShallowEq::False;
            }
        }
        return ShallowEq::RecurseAttrs(a, b);
    }
    ShallowEq::False
}

fn bool_outcome<'gc>(b: bool) -> ShallowEq<'gc> {
    if b { ShallowEq::True } else { ShallowEq::False }
}