use std::cell::RefCell;

use nixjit_error::Result;
use nixjit_hir::{Downgrade, DowngradeContext, Hir};
use nixjit_ir::ExprId;

use super::Context;

pub struct DowngradeCtx<'ctx, 'bump> {
    ctx: &'ctx mut Context<'bump>,
    irs: Vec<RefCell<Hir>>,
}

impl<'ctx, 'bump> DowngradeCtx<'ctx, 'bump> {
    pub fn new(ctx: &'ctx mut Context<'bump>) -> Self {
        Self {
            ctx,
            irs: Vec::new(),
        }
    }
}

impl DowngradeCtx<'_, '_> {
    fn get_ir(&self, id: ExprId) -> &RefCell<Hir> {
        // SAFETY: The `ExprId` is guaranteed to be valid and correspond to an expression
        // allocated within this context, making the raw index access safe.
        let idx = unsafe { id.raw() } - self.ctx.lirs.len() - self.ctx.hirs.len();
        if cfg!(debug_assertions) {
            self.irs.get(idx).unwrap()
        } else {
            // SAFETY: The index calculation is guarded by the logic that creates `ExprId`s,
            // ensuring it's always within the bounds of the `irs` vector in release builds.
            // The debug build's `unwrap()` serves as a runtime check for this invariant.
            unsafe { self.irs.get_unchecked(idx) }
        }
    }
}

impl DowngradeContext for DowngradeCtx<'_, '_> {
    fn new_expr(&mut self, expr: Hir) -> ExprId {
        self.irs.push(expr.into());
        self.ctx.alloc_id()
    }

    fn with_expr_mut<T>(&mut self, id: ExprId, f: impl FnOnce(&mut Hir, &mut Self) -> T) -> T {
        // SAFETY: This is a common pattern to temporarily bypass the borrow checker.
        // We are creating a mutable reference to `self` from a raw pointer. This is safe
        // because `self_mut` is only used within the closure `f`, and we are careful
        // not to create aliasing mutable references. The `RefCell`'s runtime borrow
        // checking further ensures that we don't have multiple mutable borrows of the
        // same `Hir` expression simultaneously.
        unsafe {
            let self_mut = &mut *(self as *mut Self);
            f(&mut self.get_ir(id).borrow_mut(), self_mut)
        }
    }

    fn downgrade_root(mut self, root: rnix::ast::Expr) -> Result<ExprId> {
        let id = root.downgrade(&mut self)?;
        self.ctx
            .hirs
            .extend(self.irs.into_iter().map(RefCell::into_inner));
        for (idx, ir) in self.ctx.hirs.iter().enumerate() {
            println!(
                "{:?} {:#?}",
                // SAFETY: The index `idx` is obtained from iterating over `self.ctx.hirs`,
                // so it is guaranteed to be a valid index. The length of `lirs` is added
                // as an offset to ensure the `ExprId` correctly corresponds to its position
                // in the combined IR storage.
                unsafe { ExprId::from_raw(idx + self.ctx.lirs.len()) },
                &ir
            );
        }
        Ok(id)
    }
}