#![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_bytecode::{InstructionPtr, PrimOpPhase};
use fix_error::{Error, Result, Source};
use fix_lang::{BUILTINS, BuiltinId, StringId};
use gc_arena::metrics::Pacing;
use gc_arena::{Arena, Collect, Gc, Mutation, RefLock, Rootable};
use hashbrown::HashMap;
use smallvec::SmallVec;

#[cfg(feature = "tailcall")]
mod dispatch_tailcall;
pub use fix_runtime::*;
mod instructions;
mod primops;

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

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

    env: GcEnv<'gc>,

    import_cache: HashMap<PathBuf, Value<'gc>>,
    scope_slots: Vec<Value<'gc>>,

    builtins: Value<'gc>,
    empty_list: Value<'gc>,
    empty_attrs: Value<'gc>,

    force_mode: ForceMode,

    #[collect(require_static)]
    result: Option<Result<fix_lang::Value>>,

    #[collect(require_static)]
    pending_load: Option<PendingLoad>,

    functor_sym: StringId,
}

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(idx as u8).expect("infallible");
        let name = name.strip_prefix("__").unwrap_or(name);
        let name = ctx.intern_string(name);
        let dispatch_ip = PrimOpPhase::entry_for_builtin(id).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));
    let builtins_value = Value::new_gc(builtins_set);
    *self_ref_thunk.borrow_mut(mc) =
        ThunkState::Evaluated(builtins_value.restrict().expect("builtins is not a thunk"));
    builtins_value
}

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())),

            import_cache: HashMap::new(),
            scope_slots: Vec::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,
            pending_load: None,

            functor_sym: ctx.intern_string("__functor"),
        }
    }
}

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

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

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

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

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

    #[inline(always)]
    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)]
    fn drop_n(&mut self, depth: usize) {
        self.stack.truncate(self.stack.len() - depth);
    }

    #[inline(always)]
    fn stack_len(&self) -> usize {
        self.stack.len()
    }

    #[inline(always)]
    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 } => {
                *state = ThunkState::Blackhole;
                self.call_stack.push(CallFrame {
                    thunk: Some(thunk),
                    pc: resume_pc,
                    env: self.env,
                });
                self.env = env;
                reader.set_pc(ip);
                Step::Break(Break::Force)
            }
            ThunkState::Evaluated(v) => {
                self.replace(depth, v.relax());
                Step::Continue(())
            }
            ThunkState::Apply { func, arg } => {
                self.call_stack.push(CallFrame {
                    thunk: Some(thunk),
                    pc: resume_pc,
                    env: self.env,
                });
                self.push(func);
                self.call(reader, mc, arg, resume_pc)
            }
            ThunkState::Blackhole => {
                self.finish_err(Error::eval_error("infinite recursion encountered"))
            }
        }
    }

    #[inline(always)]
    fn call(
        &mut self,
        reader: &mut BytecodeReader<'_>,
        mc: &Mutation<'gc>,
        arg: Value<'gc>,
        resume_pc: usize,
    ) -> Step {
        instructions::call(self, reader, mc, arg, resume_pc)
    }

    #[inline(always)]
    fn push_call_frame(&mut self, frame: CallFrame<'gc>) {
        self.call_stack.push(frame);
    }

    #[inline(always)]
    fn pop_call_frame(&mut self) -> Option<CallFrame<'gc>> {
        self.call_stack.pop()
    }

    #[inline(always)]
    fn call_depth(&self) -> usize {
        self.call_depth
    }

    #[inline(always)]
    fn inc_call_depth(&mut self) {
        self.call_depth += 1;
    }

    #[inline(always)]
    fn dec_call_depth(&mut self) {
        self.call_depth -= 1;
    }

    #[inline(always)]
    fn env(&self) -> GcEnv<'gc> {
        self.env
    }

    #[inline(always)]
    fn set_env(&mut self, env: GcEnv<'gc>) {
        self.env = env;
    }

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

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

    #[inline(always)]
    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)]
    fn builtins(&self) -> Value<'gc> {
        self.builtins
    }

    #[inline(always)]
    fn functor_sym(&self) -> StringId {
        self.functor_sym
    }

    #[inline(always)]
    fn empty_list(&self) -> Value<'gc> {
        self.empty_list
    }

    #[inline(always)]
    fn empty_attrs(&self) -> Value<'gc> {
        self.empty_attrs
    }

    #[inline(always)]
    fn force_mode(&self) -> ForceMode {
        self.force_mode
    }

    #[inline(always)]
    fn import_cache_get(&self, path: &std::path::Path) -> Option<Value<'gc>> {
        self.import_cache.get(path).copied()
    }

    #[inline(always)]
    fn import_cache_insert(&mut self, path: PathBuf, val: Value<'gc>) {
        self.import_cache.insert(path, val);
    }

    #[inline(always)]
    fn scope_slot(&self, idx: u32) -> Value<'gc> {
        *self
            .scope_slots
            .get(idx as usize)
            .expect("invalid scope slot")
    }

    #[inline(always)]
    fn scope_slots_push(&mut self, val: Value<'gc>) -> u32 {
        let idx = self.scope_slots.len() as u32;
        self.scope_slots.push(val);
        idx
    }

    #[inline(always)]
    fn set_pending_load(&mut self, load: PendingLoad) {
        self.pending_load = Some(load);
    }
}

enum Action {
    Continue { pc: usize },
    Done(Result<fix_lang::Value>),
    LoadFile(PendingLoad),
}

/// Compute initial heap size mirroring CppNix's strategy: 25% of physical RAM,
/// clamped to [32 MiB, 384 MiB]. Used as `Pacing::min_sleep` so the collector
/// defers the first cycle until the heap reaches this size.
fn initial_heap_size() -> usize {
    const MIN_SIZE: usize = 32 * 1024 * 1024;
    const MAX_SIZE: usize = 384 * 1024 * 1024;

    let mut sys = sysinfo::System::new();
    sys.refresh_memory();
    let total = sys.total_memory() as usize;
    let quarter = total / 4;
    quarter.clamp(MIN_SIZE, MAX_SIZE)
}

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

        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() > 0.0 {
                        arena.finish_cycle();
                    }
                }
                Action::LoadFile(load) => {
                    let source = match Source::new_file(load.path) {
                        Ok(src) => src,
                        Err(err) => break Err(Error::eval_error(format!("import failed: {err}"))),
                    };
                    let extra_scope = load.scope.map(|s| ExtraScope::ScopedImport {
                        keys: s.keys,
                        slot_id: s.slot_id,
                    });
                    let new_ip = match code.compile_with_scope(source, extra_scope, runtime) {
                        Ok(ip) => ip,
                        Err(err) => break Err(err),
                    };
                    pc = new_ip.0;
                    arena.mutate_root(|mc, root| {
                        root.env = Gc::new(mc, RefLock::new(Env::empty()));
                    });
                }
                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"))
                }
                TailResult::LoadFile => Action::LoadFile(
                    self.pending_load
                        .take()
                        .expect("TailResult::LoadFile without pending_load"),
                ),
            }
        }
    }

    #[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_bytecode::Op::*;
        use instructions::*;

        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 => op_push_smi(self, &mut reader),
                PushBigInt => op_push_bigint(self, &mut reader, mc),
                PushFloat => op_push_float(self, &mut reader),
                PushString => op_push_string(self, &mut reader),
                PushNull => op_push_null(self),
                PushTrue => op_push_true(self),
                PushFalse => op_push_false(self),

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

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

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

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

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

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

                OpNeg => op_neg(self, &mut reader, mc),
                OpNot => op_not(self, &mut reader, mc),

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

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

                Assert => op_assert(self, ctx, &mut reader, mc),

                LookupWith => op_lookup_with(self, ctx, &mut reader, mc),

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

                LoadReplBinding => op_load_repl_binding(self, &mut reader),
                LoadScopedBinding => op_load_scoped_binding(self, ctx, &mut reader, mc),

                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"));
                }
                Step::Break(Break::LoadFile) => {
                    return Action::LoadFile(
                        self.pending_load
                            .take()
                            .expect("Break::LoadFile without pending_load"),
                    );
                }
            }
        }
    }
}