use std::ptr::NonNull;

use bumpalo::{Bump, boxed::Box};
use hashbrown::HashMap;
use itertools::Itertools;
use petgraph::graphmap::DiGraphMap;

use nixjit_builtins::{
    builtins::{GLOBAL_LEN, SCOPED_LEN}, BuiltinFn, Builtins, BuiltinsContext
};
use nixjit_error::{Error, Result};
use nixjit_eval::{Args, EvalContext, Value};
use nixjit_hir::{DowngradeContext, Hir};
use nixjit_ir::{AttrSet, ExprId, Param, PrimOpId};
use nixjit_lir::Lir;

use crate::downgrade::DowngradeCtx;
use crate::eval::EvalCtx;
use crate::resolve::ResolveCtx;

mod downgrade;
mod eval;
mod resolve;

/// The main evaluation context.
///
/// This struct orchestrates the entire Nix expression evaluation process,
/// from parsing and semantic analysis to interpretation and JIT compilation.
pub struct Context<'bump> {
    ir_count: usize,
    hirs: Vec<Hir>,
    lirs: Vec<Box<'bump, Lir>>,
    /// Maps a function's body `ExprId` to its parameter definition.
    funcs: HashMap<ExprId, Param>,

    repl_scope: NonNull<HashMap<String, ExprId>>,
    global_scope: NonNull<HashMap<&'static str, ExprId>>,

    /// A dependency graph between expressions.
    graph: DiGraphMap<ExprId, ()>,

    /// A table of primitive operation implementations.
    primops: [(usize, BuiltinFn<Self>); GLOBAL_LEN + SCOPED_LEN],

    bump: &'bump Bump,
}

impl Drop for Context<'_> {
    fn drop(&mut self) {
        // SAFETY: `repl_scope` and `global_scope` are `NonNull` pointers to `HashMap`s
        // allocated within the `bump` arena. Because `NonNull` does not convey ownership,
        // Rust's drop checker will not automatically drop the pointed-to `HashMap`s when
        // the `Context` is dropped. We must manually call `drop_in_place` to ensure
        // their destructors are run. This is safe because these pointers are guaranteed
        // to be valid and non-null for the lifetime of the `Context`, as they are
        // initialized in `new()` and never deallocated or changed.
        unsafe {
            self.repl_scope.drop_in_place();
            self.global_scope.drop_in_place();
        }
    }
}

impl<'bump> Context<'bump> {
    pub fn new(bump: &'bump Bump) -> Self {
        let Builtins { global, scoped } = Builtins::new();
        let global_scope = global
            .iter()
            .enumerate()
            .map(|(idx, (k, _, _))| {
                // SAFETY: The index `idx` comes from `enumerate()` on the `global` array,
                // so it is guaranteed to be a valid, unique index for a primop LIR.
                (*k, unsafe { ExprId::from_raw(idx) })
            })
            .chain(core::iter::once((
                "builtins",
                // SAFETY: This ID corresponds to the `builtins` attrset LIR, which is
                // constructed and placed after all the global and scoped primop LIRs.
                // The index is calculated to be exactly at that position.
                unsafe { ExprId::from_raw(GLOBAL_LEN + SCOPED_LEN) },
            )))
            .collect();
        let primops = global
            .iter()
            .map(|&(_, arity, f)| (arity, f))
            .chain(scoped.iter().map(|&(_, arity, f)| (arity, f)))
            .collect_array()
            .unwrap();
        let lirs = (0..global.len())
            .map(|idx| {
                // SAFETY: The index `idx` is guaranteed to be within the bounds of the
                // `global` primops array, making it a valid raw ID for a `PrimOpId`.
                Lir::PrimOp(unsafe { PrimOpId::from_raw(idx) })
            })
            .chain((0..scoped.len()).map(|idx| {
                // SAFETY: The index `idx` is within the bounds of the `scoped` primops
                // array. Adding `GLOBAL_LEN` correctly offsets it to its position in
                // the combined `primops` table.
                Lir::PrimOp(unsafe { PrimOpId::from_raw(idx + GLOBAL_LEN) })
            }))
            .chain(core::iter::once(Lir::AttrSet(AttrSet {
                stcs: global
                    .into_iter()
                    .enumerate()
                    .map(|(idx, (name, ..))| {
                        // SAFETY: `idx` from `enumerate` is a valid index for the LIR
                        // corresponding to this global primop.
                        (name.to_string(), unsafe { ExprId::from_raw(idx) })
                    })
                    .chain(scoped.into_iter().enumerate().map(|(idx, (name, ..))| {
                        // SAFETY: `idx + GLOBAL_LEN` is a valid index for the LIR
                        // corresponding to this scoped primop.
                        (name.to_string(), unsafe {
                            ExprId::from_raw(idx + GLOBAL_LEN)
                        })
                    }))
                    .chain(core::iter::once((
                        "builtins".to_string(),
                        // SAFETY: This ID points to the `Thunk` that wraps this very
                        // `AttrSet`. The index is calculated to be one position after
                        // the `AttrSet` itself.
                        unsafe { ExprId::from_raw(GLOBAL_LEN + SCOPED_LEN + 1) },
                    )))
                    .collect(),
                ..AttrSet::default()
            })))
            .chain(core::iter::once(Lir::Thunk(
                // SAFETY: This ID points to the `builtins` `AttrSet` defined just above.
                // Its index is calculated to be at that exact position.
                unsafe { ExprId::from_raw(GLOBAL_LEN + SCOPED_LEN) },
            )))
            .map(|lir| Box::new_in(lir, bump))
            .collect_vec();
        Self {
            ir_count: lirs.len(),
            hirs: Vec::new(),
            lirs,
            funcs: HashMap::new(),

            global_scope: NonNull::from(bump.alloc(global_scope)),
            repl_scope: NonNull::from(bump.alloc(HashMap::new())),
            graph: DiGraphMap::new(),

            primops,

            bump,
        }
    }

    pub fn downgrade_ctx<'a>(&'a mut self) -> DowngradeCtx<'a, 'bump> {
        DowngradeCtx::new(self)
    }

    pub fn resolve_ctx<'a>(&'a mut self, root: ExprId) -> ResolveCtx<'a, 'bump> {
        ResolveCtx::new(self, root)
    }

    pub fn eval_ctx<'a>(&'a mut self) -> EvalCtx<'a, 'bump> {
        EvalCtx::new(self)
    }
    /// The main entry point for evaluating a Nix expression string.
    ///
    /// This function performs the following steps:
    /// 1. Parses the expression string into an `rnix` AST.
    /// 2. Downgrades the AST to the High-Level IR (HIR).
    /// 3. Resolves the HIR to the Low-Level IR (LIR).
    /// 4. Evaluates the LIR to produce a final `Value`.
    pub fn eval(&mut self, expr: &str) -> Result<nixjit_value::Value> {
        let root = rnix::Root::parse(expr);
        if !root.errors().is_empty() {
            return Err(Error::parse_error(
                root.errors().iter().map(|err| err.to_string()).join("; "),
            ));
        }
        let root = self
            .downgrade_ctx()
            .downgrade_root(root.tree().expr().unwrap())?;
        let ctx = self.resolve_ctx(root);
        ctx.resolve_root()?;
        Ok(self.eval_ctx().eval(root)?.to_public())
    }

    pub fn add_binding(&mut self, ident: &str, expr: &str) -> Result<()> {
        let root = rnix::Root::parse(expr);
        let root_expr = root
            .ok()
            .map_err(|err| Error::parse_error(err.to_string()))?
            .expr()
            .unwrap();
        let expr_id = self.downgrade_ctx().downgrade_root(root_expr)?;
        self.resolve_ctx(expr_id).resolve_root()?;
        // SAFETY: `repl_scope` is a `NonNull` pointer that is guaranteed to be valid
        // for the lifetime of `Context`. It is initialized in `new()` and the memory
        // it points to is managed by the `bump` arena. Therefore, it is safe to
        // dereference it to a mutable reference here.
        unsafe { self.repl_scope.as_mut() }.insert(ident.to_string(), expr_id);
        Ok(())
    }
}

impl Context<'_> {
    fn alloc_id(&mut self) -> ExprId {
        self.ir_count += 1;
        // SAFETY: This function is the sole source of new `ExprId`s during the
        // downgrade and resolve phases. By monotonically incrementing `ir_count`,
        // we guarantee that each ID is unique and corresponds to a valid, soon-to-be-
        // allocated slot in the IR vectors.
        unsafe { ExprId::from_raw(self.ir_count - 1) }
    }

    fn add_dep(&mut self, from: ExprId, to: ExprId) {
        self.graph.add_edge(from, to, ());
    }
}

impl BuiltinsContext for Context<'_> {}