use std::hash::Hash;
use std::marker::PhantomData;

use bumpalo::Bump;
use bumpalo::collections::Vec;
use fix_lang::{BUILTINS, BuiltinId, StringId};
use ghost_cell::{GhostCell, GhostToken};
use num_enum::TryFromPrimitive as _;
use rnix::{TextRange, ast};
use string_interner::DefaultStringInterner;

pub mod downgrade;

pub type HashMap<'ir, K, V> = hashbrown::HashMap<K, V, hashbrown::DefaultHashBuilder, &'ir Bump>;

pub type GhostIrRef<'id, 'ir> = <GhostRef<'id, 'ir> as RefExt<'ir>>::IrRef;
pub type GhostRoIrRef<'id, 'ir> = <GhostRoRef<'id, 'ir> as RefExt<'ir>>::IrRef;
pub type RawIrRef<'ir> = <RawRef<'ir> as RefExt<'ir>>::IrRef;
pub type GhostMaybeThunkRef<'id, 'ir> = <GhostRef<'id, 'ir> as RefExt<'ir>>::MaybeThunkRef;
pub type GhostRoMaybeThunkRef<'id, 'ir> = <GhostRoRef<'id, 'ir> as RefExt<'ir>>::MaybeThunkRef;

impl<'id, 'ir> Ir<'ir, GhostRoRef<'id, 'ir>> {
    /// Freeze a mutable IR reference into a read-only one, consuming the
    /// `GhostToken` to prevent any further mutation.
    pub fn freeze(this: GhostRoIrRef<'id, 'ir>, _: GhostToken<'id>) -> RawIrRef<'ir> {
        // SAFETY: The transmute is sound because:
        // - `GhostCell<'id, T>` is `#[repr(transparent)]` over `T`, so
        //   `&'ir GhostCell<'id, T>` and `&'ir T` have identical layout.
        // - `Ir<'ir, R>` is `#[repr(C)]`, and for every field that depends on
        //   `R`, instantiating `R = GhostRef<'id, 'ir>` vs `R = RawRef<'ir>`
        //   produces types of identical layout:
        //     - `R::IrRef` becomes `&'ir GhostCell<'id, Ir<...>>` vs `&'ir Ir<...>`
        //     - `R::MaybeThunkRef` becomes `&'ir GhostCell<'id, MaybeThunk>`
        //       vs `&'ir MaybeThunk`
        //     - `R::Ref<Ir<'ir, R>>` (used in `ConcatStrings::parts`) reduces
        //       to the same case as `R::IrRef`
        // - Therefore `IrRef<'id, 'ir>` and `RawIrRef<'ir>` are both
        //   pointer-sized references with the same layout.
        //
        // Consuming the `GhostToken` guarantees no `borrow_mut` calls can
        // occur afterwards, so the shared `&Ir` references reachable from a
        // `RawIrRef<'ir>` can never alias with mutable references.
        unsafe { std::mem::transmute::<GhostRoIrRef<'id, 'ir>, RawIrRef<'ir>>(this) }
    }
}

#[repr(transparent)]
pub struct GhostRoCell<'id, T: ?Sized>(GhostCell<'id, T>);

impl<'id, T> From<GhostCell<'id, T>> for GhostRoCell<'id, T> {
    fn from(value: GhostCell<'id, T>) -> Self {
        Self(value)
    }
}

impl<'id, T: ?Sized> From<&GhostCell<'id, T>> for &GhostRoCell<'id, T> {
    fn from(value: &GhostCell<'id, T>) -> Self {
        // SAFETY: `GhostRoCell` is `#[repr(transparent)]` over `GhostCell`
        // TODO: document mutability
        unsafe { std::mem::transmute(value) }
    }
}

impl<'id, T: ?Sized> From<&T> for &GhostRoCell<'id, T> {
    fn from(value: &T) -> Self {
        // SAFETY: `GhostRoCell` is `#[repr(transparent)]` over `GhostCell`,
        //         which is `#[repr(transparent)]` over `T`
        // TODO: document mutability
        unsafe { std::mem::transmute(value) }
    }
}

impl<'id, T: ?Sized> GhostRoCell<'id, T> {
    pub fn borrow<'a>(&'a self, token: &'a GhostToken<'id>) -> &'a T {
        self.0.borrow(token)
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub enum MaybeThunk {
    Int(i64),
    Float(f64),
    Bool(bool),
    Null,
    Str(StringId),
    Path(StringId),
    Thunk(ThunkId),
    Arg { layer: u8 },
    Builtin(BuiltinId),
    BuiltinConst(StringId),
    Builtins,
    ReplBinding(StringId),
    ScopedImportBinding { slot_id: u32, sym: StringId },
}

pub trait Ref<'ir> {
    type Ref<T>
    where
        T: 'ir;
}

pub trait RefExt<'ir>: Ref<'ir> {
    type Ir;
    type IrRef;
    type MaybeThunkRef;
}
impl<'ir, T: Ref<'ir> + 'ir> RefExt<'ir> for T {
    type Ir = Ir<'ir, Self>;
    type IrRef = Self::Ref<Self::Ir>;
    type MaybeThunkRef = Self::Ref<MaybeThunk>;
}

pub struct GhostRef<'id, 'ir>(PhantomData<&'ir GhostCell<'id, ()>>);
pub struct GhostRoRef<'id, 'ir>(PhantomData<&'ir GhostRoCell<'id, ()>>);
pub struct RawRef<'ir>(PhantomData<&'ir ()>);

impl<'id, 'ir> Ref<'ir> for GhostRef<'id, 'ir> {
    type Ref<T: 'ir> = &'ir GhostCell<'id, T>;
}
impl<'id, 'ir> Ref<'ir> for GhostRoRef<'id, 'ir> {
    type Ref<T: 'ir> = &'ir GhostRoCell<'id, T>;
}
impl<'ir> Ref<'ir> for RawRef<'ir> {
    type Ref<T: 'ir> = &'ir T;
}

#[repr(C)]
#[derive(Debug)]
pub enum Ir<'ir, R: RefExt<'ir> + ?Sized + 'ir> {
    Int(i64),
    Float(f64),
    Bool(bool),
    Null,
    Str(StringId),
    Path(R::IrRef),
    AttrSet {
        stcs: HashMap<'ir, StringId, (R::MaybeThunkRef, TextRange)>,
        dyns: Vec<'ir, (R::IrRef, R::MaybeThunkRef, TextRange)>,
    },
    List {
        items: Vec<'ir, R::MaybeThunkRef>,
    },
    ConcatStrings {
        parts: Vec<'ir, R::Ref<Ir<'ir, R>>>,
        force_string: bool,
    },

    // OPs
    UnOp {
        rhs: R::IrRef,
        kind: UnOpKind,
    },
    BinOp {
        lhs: R::IrRef,
        rhs: R::IrRef,
        kind: BinOpKind,
    },
    HasAttr {
        lhs: R::IrRef,
        rhs: Vec<'ir, Attr<R::IrRef>>,
    },
    Select {
        expr: R::IrRef,
        attrpath: Vec<'ir, Attr<R::IrRef>>,
        default: Option<R::IrRef>,
        span: TextRange,
    },

    // Conditionals
    If {
        cond: R::IrRef,
        consq: R::IrRef,
        alter: R::IrRef,
    },
    Assert {
        assertion: R::IrRef,
        expr: R::IrRef,
        assertion_raw: String,
        span: TextRange,
    },

    WithLookup {
        sym: StringId,
        namespaces: Vec<'ir, R::MaybeThunkRef>,
    },

    // Function related
    Func {
        body: R::IrRef,
        param: Option<Param<'ir>>,
        thunks: Vec<'ir, (ThunkId, R::IrRef)>,
    },
    Arg {
        layer: u8,
    },
    Call {
        func: R::IrRef,
        arg: R::MaybeThunkRef,
        span: TextRange,
    },

    // Builtins
    Builtins,
    Builtin(BuiltinId),
    BuiltinConst(StringId),

    // Misc
    TopLevel {
        body: R::IrRef,
        thunks: Vec<'ir, (ThunkId, R::IrRef)>,
    },
    MaybeThunk(R::MaybeThunkRef),
    ReplBinding(StringId),
    ScopedImportBinding {
        sym: StringId,
        slot_id: u32,
    },
}

#[repr(transparent)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct ThunkId(pub usize);

/// Represents a key in an attribute path.
#[allow(unused)]
#[derive(Debug)]
pub enum Attr<Ref> {
    /// A dynamic attribute key, which is an expression that must evaluate to a string.
    /// Example: `attrs.${key}`
    Dynamic(Ref, TextRange),
    /// A static attribute key.
    /// Example: `attrs.key`
    Str(StringId, TextRange),
}

/// The kinds of binary operations supported in Nix.
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
pub enum BinOpKind {
    // Arithmetic
    Add,
    Sub,
    Div,
    Mul,

    // Comparison
    Eq,
    Neq,
    Lt,
    Gt,
    Leq,
    Geq,

    // Logical
    And,
    Or,
    Impl,

    // Set/String/Path operations
    Con, // List concatenation (`++`)
    Upd, // AttrSet update (`//`)
}

/// The kinds of unary operations.
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
pub enum UnOpKind {
    Neg, // Negation (`-`)
    Not, // Logical not (`!`)
}

impl From<ast::UnaryOpKind> for UnOpKind {
    fn from(value: ast::UnaryOpKind) -> Self {
        match value {
            ast::UnaryOpKind::Invert => UnOpKind::Not,
            ast::UnaryOpKind::Negate => UnOpKind::Neg,
        }
    }
}

/// Describes the parameters of a function.
#[derive(Debug)]
pub struct Param<'ir> {
    pub required: Vec<'ir, (StringId, TextRange)>,
    pub optional: Vec<'ir, (StringId, TextRange)>,
    pub ellipsis: bool,
}

pub fn new_global_env(
    strings: &mut DefaultStringInterner,
) -> hashbrown::HashMap<StringId, MaybeThunk> {
    let mut global_env = hashbrown::HashMap::new();
    let builtins_sym = StringId(strings.get_or_intern("builtins"));
    global_env.insert(builtins_sym, MaybeThunk::Builtins);

    for (idx, &(name, _)) in BUILTINS.iter().enumerate() {
        let id = BuiltinId::try_from_primitive(idx as u8).expect("infallible");
        let name = StringId(strings.get_or_intern(name));
        global_env.insert(name, MaybeThunk::Builtin(id));
    }

    let consts = [
        (
            "__currentSystem",
            MaybeThunk::BuiltinConst(StringId(strings.get_or_intern("currentSystem"))),
        ),
        ("__langVersion", MaybeThunk::Int(6)),
        (
            "__nixVersion",
            MaybeThunk::BuiltinConst(StringId(strings.get_or_intern("nixVersion"))),
        ),
        (
            "__storeDir",
            MaybeThunk::BuiltinConst(StringId(strings.get_or_intern("storeDir"))),
        ),
        (
            "__nixPath",
            MaybeThunk::BuiltinConst(StringId(strings.get_or_intern("nixPath"))),
        ),
        ("null", MaybeThunk::Null),
        ("true", MaybeThunk::Bool(true)),
        ("false", MaybeThunk::Bool(false)),
    ];

    for (name, ir) in consts {
        let name = StringId(strings.get_or_intern(name));
        global_env.insert(name, ir);
    }

    global_env
}