#![allow(dead_code)]

use std::fmt;
use std::marker::PhantomData;
use std::mem::size_of;
use std::ops::Deref;

use fix_builtins::BuiltinId;
use fix_common::*;
use gc_arena::collect::Trace;
use gc_arena::{Collect, Gc, GcRefLock, Mutation, RefLock};
use num_enum::TryFromPrimitive;
use smallvec::SmallVec;
use string_interner::Symbol;
use string_interner::symbol::SymbolU32;

use crate::NixNum;
use crate::boxing::{RawBox, RawStore, RawTag, Value as RawValue};

mod private {
    pub trait Cealed {}
}

/// # Safety
/// 1. TAG must be unique among all implementors.
/// 2. TAG must be within 1..=7
pub(crate) unsafe trait Storable: private::Cealed {
    const TAG: (bool, u8);
}
pub(crate) trait InlineStorable: Storable + RawStore {}
pub(crate) trait GcStorable: Storable {}

macro_rules! define_value_types {
    (
        inline { $($itype:ty => $itag:expr, $iname:literal;)* }
        gc { $($gtype:ty => $gtag:expr, $gname:literal;)* }
    ) => {
        $(
        unsafe impl Storable for $itype {
            const TAG: (bool, u8) = $itag;
        }
        impl InlineStorable for $itype {}
        impl private::Cealed for $itype {}
        )*
        $(
        unsafe impl Storable for $gtype {
            const TAG: (bool, u8) = $gtag;
        }
        impl GcStorable for $gtype {}
        impl private::Cealed for $gtype {}
        )*

        const _: () = assert!(size_of::<Value<'static>>() == 8);
        $(const _: () = assert!(size_of::<$itype>() <= 6);)*
        $(const _: () = { let (_, val) = $itag; assert!(val >= 1 && val <= 7); };)*
        $(const _: () = { let (_, val) = $gtag; assert!(val >= 1 && val <= 7); };)*

        const _: () = {
            let tags: &[(bool, u8)] = &[$($itag),*, $($gtag),*];
            let mut mask_false: u8 = 0;
            let mut mask_true: u8 = 0;
            let mut i = 0;
            while i < tags.len() {
                let (neg, val) = tags[i];
                let bit = 1 << val;
                if neg {
                    assert!(mask_true & bit == 0, "duplicate true tag id");
                    mask_true |= bit;
                } else {
                    assert!(mask_false & bit == 0, "duplicate false tag id");
                    mask_false |= bit;
                }
                i += 1;
            }
        };

        unsafe impl<'gc> Collect<'gc> for Value<'gc> {
            const NEEDS_TRACE: bool = true;
            fn trace<T: Trace<'gc>>(&self, cc: &mut T) {
                let Some(tag) = self.raw.tag() else { return };
                match tag.neg_val() {
                    $(<$gtype as Storable>::TAG => unsafe {
                        self.load_gc::<$gtype>().trace(cc)
                    },)*
                    $(<$itype as Storable>::TAG => (),)*
                    (neg, val) => unreachable!("invalid tag: neg={neg}, val={val}"),
                }
            }
        }

        impl fmt::Debug for Value<'_> {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                match self.tag() {
                    None => write!(f, "Float({:?})", unsafe {
                        self.raw.float().unwrap_unchecked()
                    }),
                    $(Some(<$itype as Storable>::TAG) => write!(f, "{}({:?})", $iname, unsafe {
                        self.as_inline::<$itype>().unwrap_unchecked()
                    }),)*
                    $(Some(<$gtype as Storable>::TAG) =>
                        write!(f, "{}(..)", $gname),)*
                    Some((neg, val)) => write!(f, "Unknown(neg={neg}, val={val})"),
                }
            }
        }
    };
}

define_value_types! {
    inline {
        i32 => (false, 1), "SmallInt";
        bool => (false, 2), "Bool";
        Null => (false, 3), "Null";
        StringId => (false, 4), "SmallString";
        PrimOp => (false, 5), "PrimOp";
    }
    gc {
        i64 => (false, 6), "BigInt";
        NixString => (false, 7), "String";
        AttrSet<'_> => (true, 1), "AttrSet";
        List<'_> => (true, 2), "List";
        Thunk<'_> => (true, 3), "Thunk";
        Closure<'_> => (true, 4), "Closure";
        PrimOpApp<'_> => (true, 5), "PrimOpApp";
    }
}

/// # Nix runtime value representation
///
/// NaN-boxed value fitting in 8 bytes.
#[derive(Copy, Clone)]
#[repr(transparent)]
pub(crate) struct Value<'gc> {
    raw: RawBox,
    _marker: PhantomData<Gc<'gc, ()>>,
}

impl Default for Value<'_> {
    #[inline(always)]
    fn default() -> Self {
        Self::new_inline(Null)
    }
}

impl<'gc> Value<'gc> {
    #[inline(always)]
    fn from_raw_value(rv: RawValue) -> Self {
        Self {
            raw: RawBox::from_value(rv),
            _marker: PhantomData,
        }
    }

    /// Load a GC pointer from a value with a negative tag.
    ///
    /// # Safety
    ///
    /// The value must actually store a `Gc<'gc, T>` with the matching type.
    #[inline(always)]
    unsafe fn load_gc<T: GcStorable>(self) -> Gc<'gc, T> {
        unsafe {
            let rv = self.raw.value().unwrap_unchecked();
            let ptr: *const T = <*const T as RawStore>::from_val(rv);
            Gc::from_ptr(ptr)
        }
    }

    /// Returns the `(negative, val)` tag, or `None` for a float.
    #[inline(always)]
    fn tag(self) -> Option<(bool, u8)> {
        self.raw.tag().map(|t| t.neg_val())
    }
}

impl<'gc> Value<'gc> {
    #[inline]
    pub(crate) fn new_float(val: f64) -> Self {
        Self {
            raw: RawBox::from_float(val),
            _marker: PhantomData,
        }
    }

    #[inline]
    pub(crate) fn new_inline<T: InlineStorable>(val: T) -> Self {
        Self::from_raw_value(RawValue::store(
            unsafe { RawTag::new_unchecked(T::TAG.0, T::TAG.1) },
            val,
        ))
    }

    #[inline]
    pub(crate) fn new_gc<T: GcStorable>(gc: Gc<'gc, T>) -> Self {
        let ptr = Gc::as_ptr(gc);
        Self::from_raw_value(RawValue::store(
            unsafe { RawTag::new_unchecked(T::TAG.0, T::TAG.1) },
            ptr,
        ))
    }

    #[inline]
    pub(crate) fn make_int(val: i64, mc: &Mutation<'gc>) -> Self {
        if val >= i32::MIN as i64 && val <= i32::MAX as i64 {
            Value::new_inline(val as i32)
        } else {
            Value::new_gc(Gc::new(mc, val))
        }
    }
}

impl<'gc> Value<'gc> {
    #[inline]
    pub(crate) fn is_float(self) -> bool {
        self.raw.is_float()
    }

    #[inline]
    pub(crate) fn is<T: Storable>(self) -> bool {
        self.tag() == Some(T::TAG)
    }
}

impl<'gc> Value<'gc> {
    #[inline]
    pub(crate) fn as_float(self) -> Option<f64> {
        self.raw.float().copied()
    }

    #[inline]
    pub(crate) fn as_inline<T: InlineStorable>(self) -> Option<T> {
        if self.is::<T>() {
            Some(unsafe {
                let rv = self.raw.value().unwrap_unchecked();
                T::from_val(rv)
            })
        } else {
            None
        }
    }

    #[inline]
    pub(crate) fn as_gc<T: GcStorable>(self) -> Option<Gc<'gc, T>> {
        if self.is::<T>() {
            Some(unsafe {
                let rv = self.raw.value().unwrap_unchecked();
                let ptr: *const T = <*const T as RawStore>::from_val(rv);
                Gc::from_ptr(ptr)
            })
        } else {
            None
        }
    }

    #[inline]
    pub(crate) fn as_num(self) -> Option<NixNum> {
        if let Some(i) = self.as_inline::<i32>() {
            Some(NixNum::Int(i as i64))
        } else if let Some(gc_i) = self.as_gc::<i64>() {
            Some(NixNum::Int(*gc_i))
        } else {
            self.as_float().map(NixNum::Float)
        }
    }

    #[inline]
    pub(crate) fn restrict(self) -> Option<StrictValue<'gc>> {
        if !self.is::<Thunk<'gc>>() {
            Some(StrictValue(self))
        } else {
            None
        }
    }
}

#[derive(Copy, Clone, Default)]
#[repr(transparent)]
pub struct StaticValue(Value<'static>);

impl<'gc> From<StaticValue> for Value<'gc> {
    #[inline]
    fn from(value: StaticValue) -> Self {
        // SAFETY: StaticValue is guaranteed to not contain any `Gc`.
        unsafe { std::mem::transmute::<Value<'static>, Value<'gc>>(value.0) }
    }
}

impl StaticValue {
    #[inline]
    pub fn new_float(val: f64) -> Self {
        Self(Value::new_float(val))
    }
    #[inline]
    pub fn new_inline<T: InlineStorable>(val: T) -> Self {
        Self(Value::new_inline(val))
    }
    #[inline]
    pub fn new_primop(id: BuiltinId, arity: u8) -> Self {
        Self(Value::new_inline(PrimOp { id, arity }))
    }
    #[inline]
    pub fn is_float(self) -> bool {
        self.0.is_float()
    }
    #[inline]
    pub fn is<T: InlineStorable>(self) -> bool {
        self.0.is::<T>()
    }
    #[inline]
    pub fn as_float(self) -> Option<f64> {
        self.0.as_float()
    }
    #[inline]
    pub fn as_inline<T: InlineStorable>(self) -> Option<T> {
        self.0.as_inline::<T>()
    }
    #[inline]
    pub fn to_bits(self) -> u64 {
        self.0.raw.to_bits()
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct Null;
impl RawStore for Null {
    fn to_val(self, value: &mut RawValue) {
        value.set_data([0; 6]);
    }
    fn from_val(_: &RawValue) -> Self {
        Self
    }
}

impl RawStore for StringId {
    fn to_val(self, value: &mut RawValue) {
        (self.0.to_usize() as u32).to_val(value);
    }
    fn from_val(value: &RawValue) -> Self {
        Self(
            SymbolU32::try_from_usize(u32::from_val(value) as usize)
                .expect("failed to read StringId from Value"),
        )
    }
}

/// Heap-allocated Nix string.
///
/// Stored on the GC heap via `Gc<'gc, NixString>`. The string data itself
/// lives in a standard `Box<str>` owned by this struct; the GC only manages
/// the outer allocation.
#[derive(Collect)]
#[collect(require_static)]
pub(crate) struct NixString {
    data: Box<str>,
    // TODO: string context for derivation dependency tracking
}

impl NixString {
    pub(crate) fn new(s: impl Into<Box<str>>) -> Self {
        Self { data: s.into() }
    }

    pub(crate) fn as_str(&self) -> &str {
        &self.data
    }
}

impl fmt::Debug for NixString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(&self.data, f)
    }
}

#[derive(Collect, Debug, Default)]
#[collect(no_drop)]
pub(crate) struct AttrSet<'gc> {
    entries: SmallVec<[(StringId, Value<'gc>); 4]>,
}

impl<'gc> Deref for AttrSet<'gc> {
    type Target = [(StringId, Value<'gc>)];
    fn deref(&self) -> &Self::Target {
        &self.entries
    }
}

impl<'gc> AttrSet<'gc> {
    pub(crate) fn from_sorted_unchecked(entries: SmallVec<[(StringId, Value<'gc>); 4]>) -> Self {
        debug_assert!(entries.is_sorted_by_key(|(key, _)| *key));
        Self { entries }
    }

    pub(crate) fn lookup(&self, key: StringId) -> Option<Value<'gc>> {
        self.entries
            .binary_search_by_key(&key, |(k, _)| *k)
            .ok()
            .map(|i| self.entries[i].1)
    }

    pub(crate) fn has(&self, key: StringId) -> bool {
        self.entries.binary_search_by_key(&key, |(k, _)| *k).is_ok()
    }

    pub(crate) fn merge(&self, other: &Self, mc: &Mutation<'gc>) -> Gc<'gc, Self> {
        use std::cmp::Ordering::*;

        debug_assert!(self.entries.is_sorted_by_key(|(key, _)| *key));
        debug_assert!(other.entries.is_sorted_by_key(|(key, _)| *key));

        let mut entries = SmallVec::new();
        let mut i = 0;
        let mut j = 0;
        while i < self.entries.len() && j < other.entries.len() {
            match self.entries[i].0.cmp(&other.entries[j].0) {
                Less => {
                    entries.push(self.entries[i]);
                    i += 1;
                }
                Greater => {
                    entries.push(other.entries[j]);
                    j += 1;
                }
                Equal => {
                    entries.push(other.entries[j]);
                    i += 1;
                    j += 1;
                }
            }
        }
        entries.extend(other.entries[j..].iter().cloned());
        entries.extend(self.entries[i..].iter().cloned());

        debug_assert!(entries.is_sorted_by_key(|(key, _)| *key));

        Gc::new(mc, AttrSet { entries })
    }
}

#[derive(Collect, Debug, Default)]
#[collect(no_drop)]
pub(crate) struct List<'gc> {
    pub(crate) inner: SmallVec<[Value<'gc>; 4]>,
}
impl<'gc> Deref for List<'gc> {
    type Target = SmallVec<[Value<'gc>; 4]>;
    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

pub(crate) type Thunk<'gc> = RefLock<ThunkState<'gc>>;

#[derive(Collect, Debug)]
#[collect(no_drop)]
pub(crate) enum ThunkState<'gc> {
    Pending {
        ip: usize,
        env: GcEnv<'gc>,
        with_env: Option<GcWithEnv<'gc>>,
    },
    Apply {
        func: Value<'gc>,
        arg: Value<'gc>,
    },
    Blackhole,
    Evaluated(Value<'gc>),
}

#[derive(Collect, Debug)]
#[collect(no_drop)]
pub(crate) struct Env<'gc> {
    pub(crate) locals: SmallVec<[Value<'gc>; 4]>,
    pub(crate) prev: Option<GcEnv<'gc>>,
}
pub(crate) type GcEnv<'gc> = GcRefLock<'gc, Env<'gc>>;

#[derive(Collect, Debug)]
#[collect(no_drop)]
pub(crate) struct WithEnv<'gc> {
    pub(crate) env: Value<'gc>,
    pub(crate) prev: Option<GcWithEnv<'gc>>,
}
pub(crate) type GcWithEnv<'gc> = Gc<'gc, WithEnv<'gc>>;

impl<'gc> Env<'gc> {
    pub(crate) fn empty() -> Self {
        Env {
            locals: SmallVec::new(),
            prev: None,
        }
    }

    pub(crate) fn with_arg(
        arg: Value<'gc>,
        n_locals: u32,
        prev: Gc<'gc, RefLock<Env<'gc>>>,
    ) -> Self {
        let mut locals = smallvec::smallvec![Value::default(); 1 + n_locals as usize];
        locals[0] = arg;
        Env {
            locals,
            prev: Some(prev),
        }
    }
}

#[derive(Collect, Debug)]
#[collect(no_drop)]
pub(crate) struct Closure<'gc> {
    pub(crate) ip: u32,
    pub(crate) n_locals: u32,
    pub(crate) env: Gc<'gc, RefLock<Env<'gc>>>,
    pub(crate) pattern: Option<Gc<'gc, PatternInfo>>,
}

#[derive(Collect, Debug)]
#[collect(require_static)]
pub(crate) struct PatternInfo {
    pub(crate) required: SmallVec<[StringId; 4]>,
    pub(crate) optional: SmallVec<[StringId; 4]>,
    pub(crate) ellipsis: bool,
    pub(crate) param_spans: Box<[(StringId, u32)]>,
}

#[derive(Clone, Copy, Debug, Collect)]
#[collect(require_static)]
pub(crate) struct PrimOp {
    pub(crate) id: BuiltinId,
    pub(crate) arity: u8,
}

impl RawStore for PrimOp {
    fn to_val(self, value: &mut RawValue) {
        value.set_data([0, 0, 0, 0, self.id as u8, self.arity]);
    }
    fn from_val(value: &RawValue) -> Self {
        let [.., id, arity] = *value.data();
        Self {
            id: BuiltinId::try_from_primitive(id).expect("invalid BuiltinId"),
            arity,
        }
    }
}

#[derive(Collect, Debug)]
#[collect(no_drop)]
pub(crate) struct PrimOpApp<'gc> {
    pub(crate) primop: PrimOp,
    pub(crate) args: SmallVec<[Value<'gc>; 2]>,
}

#[derive(Copy, Clone, Default)]
#[repr(transparent)]
pub(crate) struct StrictValue<'gc>(Value<'gc>);

impl<'gc> StrictValue<'gc> {
    #[inline]
    pub(crate) fn relax(self) -> Value<'gc> {
        self.0
    }
}

impl<'gc> Deref for StrictValue<'gc> {
    type Target = Value<'gc>;
    #[inline]
    fn deref(&self) -> &Value<'gc> {
        &self.0
    }
}

impl fmt::Debug for StrictValue<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(&self.0, f)
    }
}

unsafe impl<'gc> Collect<'gc> for StrictValue<'gc> {
    const NEEDS_TRACE: bool = true;
    fn trace<T: gc_arena::collect::Trace<'gc>>(&self, cc: &mut T) {
        self.0.trace(cc);
    }
}