#![allow(dead_code)]

use std::cell::RefCell;
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::barrier::Unlock;
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
///
/// TAG must be unique among all implementors.
#[allow(private_interfaces)]
pub unsafe trait Storable: private::Cealed {
    const TAG: RawTag;
}
#[allow(private_bounds)]
pub trait InlineStorable: Storable + RawStore {}
pub trait GcStorable: Storable {}

macro_rules! define_value_types {
    (
        inline { $($itype:ty => $itag:expr, $iname:literal;)* }
        gc { $($gtype:ty => $gtag:expr, $gname:literal;)* }
    ) => {
        $(
        #[allow(private_interfaces)]
        unsafe impl Storable for $itype {
            const TAG: RawTag = $itag;
        }
        impl InlineStorable for $itype {}
        impl private::Cealed for $itype {}
        )*
        $(
        #[allow(private_interfaces)]
        unsafe impl Storable for $gtype {
            const TAG: RawTag = $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 tags: &[(bool, u8)] = &[$(RawTag::neg_val($itag)),*, $(RawTag::neg_val($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 {
                    $(<$gtype as Storable>::TAG => unsafe {
                        self.load_gc::<$gtype>().trace(cc)
                    },)*
                    $(<$itype as Storable>::TAG => (),)*
                    _ => unreachable!("invalid value tag"),
                }
            }
        }

        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),)*
                    _ => unreachable!("invalid value tag"),
                }
            }
        }
    };
}

define_value_types! {
    inline {
        i32 => RawTag::P1, "SmallInt";
        bool => RawTag::P2, "Bool";
        Null => RawTag::P3, "Null";
        StringId => RawTag::P4, "SmallString";
        PrimOp => RawTag::P5, "PrimOp";
    }
    gc {
        i64 => RawTag::P6, "BigInt";
        NixString => RawTag::P7, "String";
        AttrSet<'_> => RawTag::N1, "AttrSet";
        List<'_> => RawTag::N2, "List";
        Thunk<'_> => RawTag::N3, "Thunk";
        Closure<'_> => RawTag::N4, "Closure";
        PrimOpApp<'_> => RawTag::N5, "PrimOpApp";
    }
}

/// # Nix runtime value representation
///
/// NaN-boxed value fitting in 8 bytes.
#[derive(Copy, Clone)]
#[repr(transparent)]
pub 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)
        }
    }

    #[inline(always)]
    const fn tag(self) -> Option<RawTag> {
        self.raw.tag()
    }
}

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

    #[inline]
    pub fn new_inline<T: InlineStorable>(val: T) -> Self {
        Self::from_raw_value(RawValue::store(T::TAG, val))
    }

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

    #[inline]
    pub 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 fn is_float(self) -> bool {
        self.raw.is_float()
    }

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

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

    #[inline]
    pub 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 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 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 fn restrict(self) -> Result<StrictValue<'gc>, Gc<'gc, Thunk<'gc>>> {
        if let Some(thunk) = self.as_gc::<Thunk<'gc>>() {
            Err(thunk)
        } else {
            Ok(StrictValue(self))
        }
    }

    #[inline]
    pub fn ty(self) -> NixType {
        if self.is_float() {
            NixType::Float
        } else if self.is::<i32>() || self.is::<i64>() {
            NixType::Int
        } else if self.is::<bool>() {
            NixType::Bool
        } else if self.is::<Null>() {
            NixType::Null
        } else if self.is::<StringId>() {
            NixType::String
        } else if self.is::<PrimOp>() {
            NixType::PrimOp
        } else if self.is::<NixString>() {
            NixType::String
        } else if self.is::<AttrSet>() {
            NixType::AttrSet
        } else if self.is::<List>() {
            NixType::List
        } else if self.is::<Thunk>() {
            NixType::Thunk
        } else if self.is::<Closure>() {
            NixType::Closure
        } else if self.is::<PrimOpApp>() {
            NixType::PrimOpApp
        } else {
            unreachable!("value has no recognized type tag")
        }
    }

    #[inline]
    pub(crate) fn expect_inline<T: InlineStorable>(self) -> Result<T, NixType> {
        self.as_inline::<T>().ok_or_else(|| self.ty())
    }

    #[inline]
    pub(crate) fn expect_gc<T: GcStorable>(self) -> Result<Gc<'gc, T>, NixType> {
        self.as_gc::<T>().ok_or_else(|| self.ty())
    }

    #[inline]
    pub(crate) fn expect_num(self) -> Result<NixNum, NixType> {
        self.as_num().ok_or_else(|| self.ty())
    }

    #[inline]
    pub(crate) fn expect_bool(self) -> Result<bool, NixType> {
        self.as_inline::<bool>().ok_or_else(|| self.ty())
    }

    #[inline]
    pub(crate) fn expect_float(self) -> Result<f64, NixType> {
        self.as_float().ok_or_else(|| self.ty())
    }
}

#[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, dispatch_ip: u32) -> Self {
        Self(Value::new_inline(PrimOp {
            id,
            arity,
            dispatch_ip,
        }))
    }
    #[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)]
#[repr(transparent)]
#[collect(no_drop)]
pub(crate) struct List<'gc> {
    pub(crate) inner: RefLock<SmallVec<[Value<'gc>; 4]>>,
}

impl<'gc> List<'gc> {
    pub(crate) fn new_gc(mc: &Mutation<'gc>) -> Gc<'gc, Self> {
        Gc::new(mc, Self::default())
    }
}

impl<'gc> Unlock for List<'gc> {
    type Unlocked = RefCell<SmallVec<[Value<'gc>; 4]>>;
    unsafe fn unlock_unchecked(&self) -> &Self::Unlocked {
        unsafe { self.inner.unlock_unchecked() }
    }
}

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(StrictValue<'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)]>,
}

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

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

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

#[derive(Copy, Clone, Default, Collect)]
#[repr(transparent)]
#[collect(no_drop)]
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)
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Collect)]
#[collect(require_static)]
pub(crate) enum NixType {
    Int,
    Float,
    Bool,
    Null,
    String,
    AttrSet,
    List,
    Thunk,
    Closure,
    PrimOp,
    PrimOpApp,
}

impl NixType {
    pub fn display(self) -> &'static str {
        use NixType::*;
        match self {
            Int => "an integer",
            Float => "a float",
            Bool => "a boolean",
            Null => "null",
            String => "a string",
            AttrSet => "a set",
            List => "a list",
            Thunk => "a thunk",
            Closure => "a function",
            PrimOp => "a built-in function",
            PrimOpApp => "a partially applied built-in function",
        }
    }
}

impl std::fmt::Display for NixType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.display())
    }
}