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

use fix_lang::*;
use gc_arena::barrier::Unlock;
use gc_arena::collect::Trace;
use gc_arena::{Collect, Gc, GcRefLock, Mutation, RefLock};
use smallvec::SmallVec;
use string_interner::Symbol;
use string_interner::symbol::SymbolU32;

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

mod private {
    pub trait Cealed {}
}

pub trait ValueVariant: private::Cealed {
    type Ty<'gc>: 'gc;
    const TYPE: NixType;
    fn is_value(value: &Value<'_>) -> bool;
    /// # Safety
    ///
    /// [`Self::is_value`] must hold for `value`.
    unsafe fn from_raw<'gc>(value: &Value<'gc>) -> Self::Ty<'gc>;
}

/// # Safety
///
/// Each implementor must round-trip through [`Self::to_raw_box`] /
/// [`Self::from_raw_box`] and must be the sole owner of its NaN-boxed
/// representation (tagged payload or float).
pub(crate) unsafe trait Storable: private::Cealed {
    fn to_raw_box(self) -> RawBox;
    /// # Safety
    ///
    /// `raw` must represent a valid `Self`.
    unsafe fn from_raw_box(raw: &RawBox) -> Self;
}

macro_rules! define_value_types {
    (
        inline { $($itype:ty => $itag:path, $ity:path, $iname:literal;)* }
        gc { $($gtype:ty => $gtag:path, $gty:path, $gname:literal;)* }
    ) => {
        $(
        unsafe impl Storable for $itype {
            fn to_raw_box(self) -> RawBox {
                RawBox::from_value(RawValue::store($itag, self))
            }
            unsafe fn from_raw_box(raw: &RawBox) -> Self {
                unsafe { <Self as RawStore>::from_val(raw.value().unwrap_unchecked()) }
            }
        }
        impl private::Cealed for $itype {}
        impl ValueVariant for $itype {
            type Ty<'gc> = $itype;
            const TYPE: NixType = $ity;
            #[inline(always)]
            fn is_value(value: &Value<'_>) -> bool {
                value.raw.tag() == Some($itag)
            }
            #[inline(always)]
            unsafe fn from_raw<'gc>(value: &Value<'gc>) -> Self::Ty<'gc> {
                unsafe { <$itype as Storable>::from_raw_box(&value.raw) }
            }
        }
        )*
        $(
        unsafe impl Storable for Gc<'_, $gtype> {
            fn to_raw_box(self) -> RawBox {
                RawBox::from_value(RawValue::store($gtag, Gc::as_ptr(self)))
            }
            unsafe fn from_raw_box(raw: &RawBox) -> Self {
                unsafe { Gc::from_ptr(<*mut $gtype as RawStore>::from_val(raw.value().unwrap_unchecked())) }
            }
        }
        impl private::Cealed for Gc<'_, $gtype> {}
        impl private::Cealed for $gtype {}
        impl ValueVariant for $gtype {
            type Ty<'gc> = Gc<'gc, fix_macros::unelide_lifetimes!('gc; $gtype)>;
            const TYPE: NixType = $gty;
            #[inline(always)]
            fn is_value(value: &Value<'_>) -> bool {
                value.raw.tag() == Some($gtag)
            }
            #[inline(always)]
            unsafe fn from_raw<'gc>(value: &Value<'gc>) -> Self::Ty<'gc> {
                unsafe {
                    <Gc<'gc, fix_macros::unelide_lifetimes!('gc; $gtype)> as Storable>::from_raw_box(
                        &value.raw,
                    )
                }
            }
        }
        )*

        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 {
                    $($gtag => unsafe {
                        self.downcast::<$gtype>().unwrap_unchecked().trace(cc)
                    },)*
                    $($itag => (),)*
                    _ => 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($itag) => write!(f, "{}({:?})", $iname, unsafe {
                        self.downcast::<$itype>().unwrap_unchecked()
                    }),)*
                    $(Some($gtag) =>
                        write!(f, "{}(..)", $gname),)*
                    _ => unreachable!("invalid value tag"),
                }
            }
        }
    };
}

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

impl private::Cealed for f64 {}

unsafe impl Storable for f64 {
    fn to_raw_box(self) -> RawBox {
        RawBox::from_float(self)
    }
    unsafe fn from_raw_box(raw: &RawBox) -> Self {
        unsafe { raw.float().copied().unwrap_unchecked() }
    }
}

impl ValueVariant for f64 {
    type Ty<'gc> = f64;
    const TYPE: NixType = NixType::Float;
    #[inline(always)]
    fn is_value(value: &Value<'_>) -> bool {
        value.raw.is_float()
    }
    #[inline(always)]
    unsafe fn from_raw<'gc>(value: &Value<'gc>) -> Self::Ty<'gc> {
        unsafe { <f64 as Storable>::from_raw_box(&value.raw) }
    }
}

/// # 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(Null)
    }
}

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

impl<'gc> Value<'gc> {
    #[inline]
    #[allow(private_bounds)]
    pub fn new<T: Storable>(val: T) -> Self {
        Self {
            raw: val.to_raw_box(),
            _marker: PhantomData,
        }
    }

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

    #[inline]
    pub fn is<T: ValueVariant>(self) -> bool {
        T::is_value(&self)
    }

    #[inline]
    pub fn downcast<T: ValueVariant>(self) -> Option<T::Ty<'gc>> {
        self.is::<T>().then(|| unsafe { T::from_raw(&self) })
    }

    #[inline]
    pub fn to_bits(self) -> u64 {
        self.raw.to_bits()
    }

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

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

    #[inline]
    pub fn ty(self) -> NixType {
        if self.is::<f64>() {
            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::<Path>() {
            NixType::Path
        } 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 fn expect<T: ValueVariant>(self) -> Result<T::Ty<'gc>, NixType> {
        self.downcast::<T>().ok_or_else(|| self.ty())
    }

    #[inline]
    pub fn expect_num(self) -> Result<NixNum, NixType> {
        self.downcast_num().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]
    #[allow(private_bounds)]
    pub fn new<T: Storable + 'static>(val: T) -> Self {
        Self(Value::new(val))
    }

    #[inline]
    pub fn is<T: ValueVariant>(self) -> bool {
        self.0.is::<T>()
    }

    #[inline]
    pub fn downcast<T: ValueVariant>(self) -> Option<T::Ty<'static>> {
        self.0.downcast::<T>()
    }

    #[inline]
    pub fn to_bits(self) -> u64 {
        self.0.raw.to_bits()
    }
}

#[derive(Clone, Copy, Debug)]
pub 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"),
        )
    }
}

/// A canonicalized absolute path. Inline value carrying an interned
/// `StringId` whose contents are the path's absolute, dot-resolved form.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Path(pub StringId);

impl RawStore for Path {
    fn to_val(self, value: &mut RawValue) {
        self.0.to_val(value);
    }
    fn from_val(value: &RawValue) -> Self {
        Self(StringId::from_val(value))
    }
}

#[derive(Collect)]
#[collect(require_static)]
pub struct NixString {
    data: Box<str>,
    context: StringContext,
}

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

    /// Construct a `NixString` whose `context` is already sorted+deduped.
    /// The caller is responsible for invariant maintenance.
    pub fn with_context(s: impl Into<Box<str>>, context: StringContext) -> Self {
        Self {
            data: s.into(),
            context,
        }
    }

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

    pub fn context(&self) -> &StringContext {
        &self.context
    }

    pub fn has_context(&self) -> bool {
        !self.context.is_empty()
    }
}

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 struct AttrSet<'gc> {
    pub entries: SmallVec<[(StringId, Value<'gc>); 4]>,
}

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

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

    pub 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 struct List<'gc> {
    pub inner: RefLock<SmallVec<[Value<'gc>; 4]>>,
}

impl<'gc> List<'gc> {
    pub fn new(mc: &Mutation<'gc>, data: SmallVec<[Value<'gc>; 4]>) -> Gc<'gc, Self> {
        Gc::new(
            mc,
            Self {
                inner: RefLock::new(data),
            },
        )
    }

    pub 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 type Thunk<'gc> = RefLock<ThunkState<'gc>>;

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

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

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

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

    pub 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 struct Closure<'gc> {
    pub ip: u32,
    pub n_locals: u32,
    pub env: Gc<'gc, RefLock<Env<'gc>>>,
    pub pattern: Option<Gc<'gc, PatternInfo>>,
}

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

#[repr(packed, Rust)]
#[derive(Clone, Copy, Debug, Collect)]
#[collect(require_static)]
pub struct PrimOp {
    pub id: BuiltinId,
    pub arity: u8,
    pub 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(id).expect("invalid BuiltinId"),
            arity,
            dispatch_ip: u32::from_le_bytes(bytes),
        }
    }
}

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

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

impl<'gc> StrictValue<'gc> {
    #[inline]
    pub 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 enum NixType {
    Int,
    Float,
    Bool,
    Null,
    String,
    Path,
    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",
            Path => "a path",
            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())
    }
}

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