#![warn(clippy::unwrap_used)]
#![cfg_attr(feature = "tailcall", expect(incomplete_features))]
#![cfg_attr(
    feature = "tailcall",
    feature(explicit_tail_calls, rust_preserve_none_cc)
)]

use std::path::PathBuf;

use fix_builtins::{BUILTINS, BuiltinId};
use fix_codegen::InstructionPtr;
use fix_common::StringId;
use fix_error::{Error, Result, Source};
use gc_arena::arena::CollectionPhase;
use gc_arena::{Arena, Collect, Gc, Mutation, RefLock, Rootable};
use hashbrown::HashMap;
use num_enum::TryFromPrimitive;
use smallvec::SmallVec;

mod boxing;
mod bytecode_reader;
#[cfg(feature = "tailcall")]
mod dispatch_tailcall;
mod forced;
mod value;
pub use value::StaticValue;
use value::*;
mod helpers;
pub(crate) mod instructions;
pub(crate) use bytecode_reader::BytecodeReader;
pub(crate) use forced::Forced;
use helpers::*;

type VmResult<T> = std::result::Result<T, VmError>;

#[allow(dead_code)]
enum VmError {
    Catchable(String),
    Uncatchable(Box<Error>),
}

impl From<Box<Error>> for VmError {
    fn from(e: Box<Error>) -> Self {
        VmError::Uncatchable(e)
    }
}

impl VmError {
    fn into_error(self) -> Box<Error> {
        match self {
            VmError::Catchable(_) => todo!("Check for tryEval catch frames"),
            VmError::Uncatchable(e) => e,
        }
    }
}

#[derive(Collect, Clone, Copy, Debug, PartialEq, Eq, Default)]
#[collect(require_static)]
pub enum ForceMode {
    #[default]
    AsIs,
    Shallow,
    Deep,
}

pub trait VmContext {
    fn split(&mut self) -> (&mut impl VmCode, &mut impl VmRuntimeCtx);
}

pub trait VmRuntimeCtx {
    fn intern_string(&mut self, s: impl AsRef<str>) -> StringId;
    fn resolve_string(&self, id: StringId) -> &str;
    fn get_const(&self, id: u32) -> StaticValue;
    fn add_const(&mut self, val: StaticValue) -> u32;
}

pub trait VmCode {
    fn bytecode(&self) -> &[u8];
    fn compile(
        &mut self,
        source: Source,
        ctx: &mut impl VmRuntimeCtx,
    ) -> fix_error::Result<InstructionPtr>;
}

pub(crate) trait VmRuntimeCtxExt: VmRuntimeCtx {
    fn get_string<'a, 'gc: 'a>(&'a self, val: StrictValue<'gc>) -> Option<&'a str>;
    fn get_string_id<'a, 'gc: 'a>(
        &'a mut self,
        val: StrictValue<'gc>,
    ) -> std::result::Result<StringId, NixType>;
    fn convert_value(&self, val: Value) -> fix_common::Value;
}

impl<T: VmRuntimeCtx> VmRuntimeCtxExt for T {
    fn get_string<'a, 'gc: 'a>(&'a self, val: StrictValue<'gc>) -> Option<&'a str> {
        if let Some(sid) = val.as_inline::<StringId>() {
            Some(self.resolve_string(sid))
        } else {
            val.as_gc::<NixString>().map(|ns| ns.as_ref().as_str())
        }
    }

    fn get_string_id<'a, 'gc: 'a>(
        &'a mut self,
        val: StrictValue<'gc>,
    ) -> std::result::Result<StringId, NixType> {
        if let Some(sid) = val.as_inline::<StringId>() {
            Ok(sid)
        } else if let Some(s) = val.as_gc::<NixString>().map(|ns| ns.as_ref().as_str()) {
            Ok(self.intern_string(s))
        } else {
            Err(val.ty())
        }
    }

    fn convert_value(&self, val: Value) -> fix_common::Value {
        use fix_common::Value;
        if let Some(i) = val.as_inline::<i32>() {
            Value::Int(i as i64)
        } else if let Some(gc_i) = val.as_gc::<i64>() {
            Value::Int(*gc_i)
        } else if let Some(f) = val.as_float() {
            Value::Float(f)
        } else if let Some(b) = val.as_inline::<bool>() {
            Value::Bool(b)
        } else if val.is::<Null>() {
            Value::Null
        } else if let Some(sid) = val.as_inline::<StringId>() {
            let s = self.resolve_string(sid).to_owned();
            Value::String(s)
        } else if let Some(ns) = val.as_gc::<NixString>() {
            Value::String(ns.as_str().to_owned())
        } else if let Some(attrs) = val.as_gc::<AttrSet>() {
            let mut map = std::collections::BTreeMap::new();
            for &(key, val) in attrs.iter() {
                let key = self.resolve_string(key).to_owned();
                let converted = self.convert_value(val);
                map.insert(fix_common::Symbol::from(key), converted);
            }
            Value::AttrSet(fix_common::AttrSet::new(map))
        } else if let Some(list) = val.as_gc::<List>() {
            let items: Vec<_> = list
                .inner
                .borrow()
                .iter()
                .copied()
                .map(|v| self.convert_value(v))
                .collect();
            Value::List(fix_common::List::new(items))
        } else if val.is::<Closure>() {
            Value::Func
        } else if val.is::<Thunk>() {
            Value::Thunk
        } else if let Some(primop) = val.as_inline::<PrimOp>() {
            let name = fix_builtins::BUILTINS[primop.id as usize].0;
            Value::PrimOp(name.strip_prefix("__").unwrap_or(name))
        } else if let Some(app) = val.as_gc::<PrimOpApp>() {
            let name = fix_builtins::BUILTINS[app.primop.id as usize].0;
            Value::PrimOpApp(name.strip_prefix("__").unwrap_or(name))
        } else {
            Value::Null
        }
    }
}

#[repr(u8)]
pub(crate) enum Break {
    Force,
    Done,
}

pub(crate) type Step = std::ops::ControlFlow<Break>;

#[derive(Collect)]
#[collect(no_drop)]
pub struct Vm<'gc> {
    pub(crate) stack: Vec<Value<'gc>>,
    pub(crate) call_stack: Vec<CallFrame<'gc>>,
    pub(crate) call_depth: usize,
    #[allow(dead_code)]
    #[collect(require_static)]
    pub(crate) error_context: Vec<ErrorFrame>,

    pub(crate) env: GcEnv<'gc>,
    pub(crate) with_env: Option<GcWithEnv<'gc>>,

    pub(crate) import_cache: HashMap<PathBuf, Value<'gc>>,

    pub(crate) builtins: Value<'gc>,
    pub(crate) empty_list: Value<'gc>,
    pub(crate) empty_attrs: Value<'gc>,

    pub(crate) force_mode: ForceMode,

    #[collect(require_static)]
    pub(crate) result: Option<Result<fix_common::Value>>,
}

pub(crate) enum OperandData {
    Const(StaticValue),
    Local { layer: u8, idx: u32 },
    Builtins,
    BigInt(i64),
}

impl OperandData {
    pub(crate) fn resolve<'gc>(&self, mc: &Mutation<'gc>, root: &Vm<'gc>) -> Value<'gc> {
        match *self {
            OperandData::Const(sv) => sv.into(),
            OperandData::Local { layer, idx } => {
                let mut cur = root.env;
                for _ in 0..layer {
                    let prev = cur.borrow().prev.expect("env chain too short");
                    cur = prev;
                }
                cur.borrow().locals[idx as usize]
            }
            OperandData::Builtins => root.builtins,
            OperandData::BigInt(val) => Value::new_gc(Gc::new(mc, val)),
        }
    }
}

pub(crate) enum AttrKeyData {
    Static(StringId),
    Dynamic(OperandData),
}

fn init_builtins<'gc>(mc: &Mutation<'gc>, ctx: &mut impl VmRuntimeCtx) -> Value<'gc> {
    let mut entries = SmallVec::with_capacity(BUILTINS.len());

    for (idx, &(name, arity)) in BUILTINS.iter().enumerate() {
        let id = BuiltinId::try_from_primitive(idx as u8).expect("infallible");
        let name = name.strip_prefix("__").unwrap_or(name);
        let name = ctx.intern_string(name);
        let dispatch_ip = id.entry_phase().ip();
        entries.push((
            name,
            Value::new_inline(PrimOp {
                id,
                arity,
                dispatch_ip,
            }),
        ));
    }

    let consts = [
        (
            "__currentSystem",
            Value::new_inline(ctx.intern_string("x86_64-linux")),
        ),
        ("__langVersion", Value::new_inline(6i32)),
        (
            "__nixVersion",
            Value::new_inline(ctx.intern_string("2.24.0")),
        ),
        (
            "__storeDir",
            Value::new_inline(ctx.intern_string("/nix/store")),
        ),
        ("__nixPath", Value::new_gc(Gc::new(mc, List::default()))),
        ("null", Value::new_inline(Null)),
        ("true", Value::new_inline(true)),
        ("false", Value::new_inline(false)),
    ];

    for (name, val) in consts {
        let name = name.strip_prefix("__").unwrap_or(name);
        let name = ctx.intern_string(name);
        entries.push((name, val));
    }

    let self_ref_thunk = Gc::new(mc, RefLock::new(ThunkState::Blackhole));
    let sym = ctx.intern_string("builtins");
    entries.push((sym, Value::new_gc(self_ref_thunk)));

    entries.sort_by_key(|(k, _)| *k);

    let builtins_set = Gc::new(mc, AttrSet::from_sorted_unchecked(entries));
    Value::new_gc(builtins_set)
}

impl<'gc> Vm<'gc> {
    fn new(force_mode: ForceMode, mc: &Mutation<'gc>, ctx: &mut impl VmRuntimeCtx) -> Self {
        let builtins = init_builtins(mc, ctx);
        Vm {
            stack: Vec::with_capacity(8192),
            call_stack: Vec::with_capacity(1024),
            call_depth: 0,
            error_context: Vec::with_capacity(1024),

            env: Gc::new(mc, RefLock::new(Env::empty())),
            with_env: None,

            import_cache: HashMap::new(),

            builtins,
            empty_list: Value::new_gc(Gc::new(mc, List::default())),
            empty_attrs: Value::new_gc(Gc::new(mc, AttrSet::default())),

            force_mode,

            result: None,
        }
    }

    #[inline(always)]
    pub(crate) fn finish_ok(&mut self, val: fix_common::Value) -> Step {
        self.result = Some(Ok(val));
        Step::Break(Break::Done)
    }

    #[inline(always)]
    pub(crate) fn finish_err(&mut self, err: Box<Error>) -> Step {
        self.result = Some(Err(err));
        Step::Break(Break::Done)
    }

    #[inline(always)]
    pub(crate) fn finish_type_err(&mut self, expected: NixType, got: NixType) -> Step {
        self.result = Some(Err(Error::eval_error(format!(
            "expected {expected}, got {got}"
        ))));
        Step::Break(Break::Done)
    }

    #[inline(always)]
    pub(crate) fn finish_vm_err(&mut self, err: VmError) -> Step {
        self.finish_err(err.into_error())
    }

    #[inline(always)]
    pub(crate) fn push(&mut self, val: Value<'gc>) {
        self.stack.push(val);
    }

    #[inline(always)]
    #[must_use]
    pub(crate) fn pop(&mut self) -> Value<'gc> {
        self.stack.pop().expect("stack underflow")
    }

    #[inline(always)]
    #[must_use]
    pub(crate) fn peek(&mut self, depth: usize) -> Value<'gc> {
        *self
            .stack
            .get(self.stack.len() - depth - 1)
            .expect("stack underflow")
    }

    #[inline(always)]
    #[must_use]
    pub(crate) fn peek_forced(&mut self, depth: usize) -> StrictValue<'gc> {
        self.stack
            .get(self.stack.len() - depth - 1)
            .expect("stack underflow")
            .restrict()
            .expect("forced")
    }

    #[inline(always)]
    pub(crate) fn replace(&mut self, depth: usize, val: Value<'gc>) {
        let len = self.stack.len();
        *self
            .stack
            .get_mut(len - depth - 1)
            .expect("stack underflow") = val;
    }

    #[inline(always)]
    #[cfg_attr(debug_assertions, track_caller)]
    pub(crate) fn pop_forced(&mut self) -> StrictValue<'gc> {
        self.stack
            .pop()
            .expect("stack underflow")
            .restrict()
            .expect("forced")
    }

    #[inline(always)]
    pub(crate) fn try_force<T: Forced<'gc>>(
        &mut self,
        reader: &mut BytecodeReader<'_>,
        mc: &Mutation<'gc>,
    ) -> std::ops::ControlFlow<Break, T> {
        self.try_force_to_pc(reader, mc, reader.inst_start_pc())
    }

    #[inline(always)]
    pub(crate) fn try_force_to_pc<T: Forced<'gc>>(
        &mut self,
        reader: &mut BytecodeReader<'_>,
        mc: &Mutation<'gc>,
        resume_pc: usize,
    ) -> std::ops::ControlFlow<Break, T> {
        T::force_and_check(self, reader, mc, 0, resume_pc)?;
        std::ops::ControlFlow::Continue(T::pop_converted(self))
    }

    #[inline(always)]
    #[allow(unused)]
    pub(crate) fn force_slot(
        &mut self,
        depth: usize,
        reader: &mut BytecodeReader<'_>,
        mc: &Mutation<'gc>,
    ) -> Step {
        self.force_slot_to_pc(depth, reader, mc, reader.inst_start_pc())
    }

    #[inline(always)]
    pub(crate) fn force_slot_to_pc(
        &mut self,
        depth: usize,
        reader: &mut BytecodeReader<'_>,
        mc: &Mutation<'gc>,
        resume_pc: usize,
    ) -> Step {
        let Some(thunk) = self.peek(depth).as_gc::<Thunk>() else {
            return Step::Continue(());
        };
        let mut state = thunk.borrow_mut(mc);
        match *state {
            ThunkState::Pending { ip, env, with_env } => {
                *state = ThunkState::Blackhole;
                drop(state);
                self.call_stack.push(CallFrame {
                    thunk: Some(thunk),
                    stack_depth: depth,
                    pc: resume_pc,
                    env: self.env,
                    with_env: self.with_env,
                });
                self.env = env;
                self.with_env = with_env;
                reader.set_pc(ip);
                Step::Break(Break::Force)
            }
            ThunkState::Evaluated(v) => {
                drop(state);
                self.replace(depth, v.relax());
                Step::Continue(())
            }
            ThunkState::Apply { .. } => todo!("force apply"),
            ThunkState::Blackhole => {
                drop(state);
                self.finish_err(Error::eval_error("infinite recursion encountered"))
            }
        }
    }
}

#[allow(dead_code)]
struct ErrorFrame {
    span_id: u32,
    message: Option<String>,
}

#[derive(Collect, Debug)]
#[collect(no_drop)]
pub(crate) struct CallFrame<'gc> {
    pub(crate) pc: usize,
    pub(crate) stack_depth: usize,
    pub(crate) thunk: Option<Gc<'gc, Thunk<'gc>>>,
    pub(crate) env: Gc<'gc, RefLock<Env<'gc>>>,
    pub(crate) with_env: Option<Gc<'gc, WithEnv<'gc>>>,
}

pub(crate) enum Action {
    Continue { pc: usize },
    Done(Result<fix_common::Value>),
}

pub(crate) enum NixNum {
    Int(i64),
    Float(f64),
}

impl Vm<'_> {
    pub fn run<C: VmContext>(
        ctx: &mut C,
        ip: InstructionPtr,
        force_mode: ForceMode,
    ) -> Result<fix_common::Value> {
        let (code, runtime) = ctx.split();
        let mut arena: Arena<Rootable![Vm<'_>]> = Arena::new(|mc| Vm::new(force_mode, mc, runtime));

        const COLLECTOR_GRANULARITY: f64 = 1024.0;

        let mut pc = ip.0;
        loop {
            let bytecode = code.bytecode();
            match arena.mutate_root(|mc, root| root.dispatch_batch(bytecode, runtime, pc, mc)) {
                Action::Continue { pc: new_pc } => {
                    pc = new_pc;
                    if arena.metrics().allocation_debt() > COLLECTOR_GRANULARITY {
                        if arena.collection_phase() == CollectionPhase::Sweeping {
                            arena.collect_debt();
                        } else if let Some(marked) = arena.mark_debt() {
                            marked.start_sweeping();
                        }
                    }
                }
                Action::Done(done) => break done,
            }
        }
    }
}

impl<'gc> Vm<'gc> {
    const DEFAULT_FUEL_AMOUNT: u32 = 1024;

    #[inline(always)]
    fn dispatch_batch<C: VmRuntimeCtx>(
        &mut self,
        bytecode: &[u8],
        ctx: &mut C,
        pc: usize,
        mc: &Mutation<'gc>,
    ) -> Action {
        #[cfg(not(feature = "tailcall"))]
        {
            self.execute_batch(bytecode, ctx, pc, mc)
        }
        #[cfg(feature = "tailcall")]
        {
            use crate::dispatch_tailcall::{TailResult, run_tailcall};
            match run_tailcall(self, mc, ctx, bytecode, pc as u32) {
                TailResult::YieldFuel(new_pc) => Action::Continue {
                    pc: new_pc as usize,
                },
                TailResult::Done => {
                    Action::Done(self.result.take().expect("TailResult::Done without result"))
                }
            }
        }
    }

    #[inline(always)]
    #[cfg(not(feature = "tailcall"))]
    fn execute_batch(
        &mut self,
        bytecode: &[u8],
        ctx: &mut impl VmRuntimeCtx,
        pc: usize,
        mc: &Mutation<'gc>,
    ) -> Action {
        use fix_codegen::Op::*;

        let mut reader = BytecodeReader::new(bytecode, pc);
        let mut fuel = Self::DEFAULT_FUEL_AMOUNT;

        loop {
            if fuel == 0 {
                return Action::Continue { pc: reader.pc() };
            }
            fuel -= 1;

            let op = reader.read_op();

            let result = match op {
                PushSmi => self.op_push_smi(&mut reader),
                PushBigInt => self.op_push_bigint(&mut reader, mc),
                PushFloat => self.op_push_float(&mut reader),
                PushString => self.op_push_string(&mut reader),
                PushNull => self.op_push_null(),
                PushTrue => self.op_push_true(),
                PushFalse => self.op_push_false(),

                LoadLocal => self.op_load_local(&mut reader),
                LoadOuter => self.op_load_outer(&mut reader),
                StoreLocal => self.op_store_local(&mut reader, mc),
                AllocLocals => self.op_alloc_locals(&mut reader, mc),

                MakeThunk => self.op_make_thunk(&mut reader, mc),
                MakeClosure => self.op_make_closure(&mut reader, mc),
                MakePatternClosure => self.op_make_pattern_closure(&mut reader, mc),

                Call => self.op_call(ctx, &mut reader, mc),
                DispatchPrimOp => self.op_dispatch_primop(ctx, &mut reader, mc),
                Return => self.op_return(ctx, &mut reader, mc),

                MakeAttrs => self.op_make_attrs(ctx, &mut reader, mc),
                MakeEmptyAttrs => self.op_make_empty_attrs(),
                SelectStatic => self.op_select_static(ctx, &mut reader, mc),
                SelectDynamic => self.op_select_dynamic(ctx, &mut reader, mc),
                HasAttrPathStatic => self.op_has_attr_path_static(ctx, &mut reader, mc),
                HasAttrPathDynamic => self.op_has_attr_path_dynamic(ctx, &mut reader, mc),
                HasAttrStatic => self.op_has_attr_static(&mut reader, mc),
                HasAttrDynamic => self.op_has_attr_dynamic(ctx, &mut reader, mc),
                HasAttrResolve => self.op_has_attr_resolve(),
                JumpIfSelectFailed => self.op_jump_if_select_failed(&mut reader),
                JumpIfSelectSucceeded => self.op_jump_if_select_succeeded(&mut reader),

                MakeList => self.op_make_list(ctx, &mut reader, mc),
                MakeEmptyList => self.op_make_empty_list(),

                OpAdd => self.op_add(ctx, &mut reader, mc),
                OpSub => self.op_sub(&mut reader, mc),
                OpMul => self.op_mul(&mut reader, mc),
                OpDiv => self.op_div(&mut reader, mc),
                OpEq => self.op_eq(ctx, &mut reader, mc),
                OpNeq => self.op_neq(ctx, &mut reader, mc),
                OpLt => self.op_lt(ctx, &mut reader, mc),
                OpGt => self.op_gt(ctx, &mut reader, mc),
                OpLeq => self.op_leq(ctx, &mut reader, mc),
                OpGeq => self.op_geq(ctx, &mut reader, mc),
                OpConcat => self.op_concat(&mut reader, mc),
                OpUpdate => self.op_update(&mut reader, mc),

                OpNeg => self.op_neg(),
                OpNot => self.op_not(),

                JumpIfFalse => self.op_jump_if_false(&mut reader, mc),
                JumpIfTrue => self.op_jump_if_true(&mut reader, mc),
                Jump => self.op_jump(&mut reader),

                ConcatStrings => self.op_concat_strings(ctx, &mut reader, mc),
                ResolvePath => self.op_resolve_path(ctx),

                Assert => self.op_assert(&mut reader),

                PushWith => self.op_push_with(ctx, &mut reader, mc),
                PopWith => self.op_pop_with(),
                LookupWith => self.op_lookup_with(ctx, &mut reader, mc),
                PrepareWith => self.op_prepare_with(),

                LoadBuiltins => self.op_load_builtins(),
                LoadBuiltin => self.op_load_builtin(&mut reader),

                MkPos => self.op_mk_pos(&mut reader),
                LoadReplBinding => self.op_load_repl_binding(&mut reader),
                LoadScopedBinding => self.op_load_scoped_binding(&mut reader),

                Illegal => unreachable!(),
            };

            match result {
                Step::Continue(()) | Step::Break(Break::Force) => {}
                Step::Break(Break::Done) => {
                    return Action::Done(self.result.take().expect("Break::Done without result"));
                }
            }
        }
    }
}