//! The high-level intermediate representation (HIR) for nixjit.
//!
//! This module defines the data structures for the HIR, which is a more abstract and
//! semantically rich representation of the original Nix code compared to the raw AST from `rnix`.
//! It's designed to be easily translatable from the AST and serves as a stepping stone
//! towards the lower-level IR (`nixjit_lir`).
//!
//! The key components are:
//! - `Hir`: An enum representing all possible expression types in the HIR.
//! - `Downgrade`: A trait for converting `rnix::ast` nodes into HIR expressions.
//! - `DowngradeContext`: A trait that provides the necessary context for the conversion,
//!   such as allocating new expressions and functions.

use derive_more::{IsVariant, TryUnwrap, Unwrap};
use hashbrown::HashMap;

use nixjit_error::{Error, Result};
use nixjit_ir::{
    Assert, Attr, AttrSet, BinOp, Call, ConcatStrings, Const, ExprId, Func, HasAttr, If, List,
    Path, Select, Str, UnOp, Var, With,
};
use nixjit_macros::ir;
use nixjit_value::format_symbol;

mod downgrade;
mod utils;
use utils::*;

pub use downgrade::Downgrade;

/// A context for the AST-to-HIR downgrading process.
///
/// This trait abstracts the storage of HIR expressions and functions, allowing the
/// `downgrade` implementations to be generic over the specific context implementation.
pub trait DowngradeContext {
    /// Allocates a new HIR expression in the context and returns its ID.
    fn new_expr(&mut self, expr: Hir) -> ExprId;

    /// Provides temporary access to an immutable expression for inspection or use.
    fn with_expr<T>(&self, id: ExprId, f: impl FnOnce(&Hir, &Self) -> T) -> T;

    /// Provides temporary mutable access to an expression.
    fn with_expr_mut<T>(&mut self, id: &ExprId, f: impl FnOnce(&mut Hir, &mut Self) -> T) -> T;
}

ir! {
    Hir,

    // Represents an attribute set, e.g., `{ a = 1; b = 2; }`.
    AttrSet,
    // Represents a list, e.g., `[1 2 3]`.
    List,
    // Represents a "has attribute" check, e.g., `attrs ? a`.
    HasAttr,
    // Represents a binary operation, e.g., `a + b`.
    BinOp,
    // Represents a unary operation, e.g., `-a`.
    UnOp,
    // Represents an attribute selection, e.g., `attrs.a` or `attrs.a or defaultValue`.
    Select,
    // Represents an if-then-else expression.
    If,
    // Represents a function definition (lambda).
    Func,
    // Represents a function call.
    Call,
    // Represents a `with` expression, e.g., `with pkgs; stdenv.mkDerivation { ... }`.
    With,
    // Represents an `assert` expression.
    Assert,
    // Represents the concatenation of strings, often from interpolated strings.
    ConcatStrings,
    // Represents a constant value (integer, float, boolean, null).
    Const,
    // Represents a simple string literal.
    Str,
    // Represents a variable lookup by its symbol/name.
    Var,
    // Represents a path expression.
    Path,
    // Represents a `let ... in ...` binding.
    Let { pub bindings: HashMap<String, ExprId>, pub body: ExprId },
    // Represents a function argument lookup.
    Arg,
}

#[derive(Debug)]
pub struct Arg;

/// A trait defining operations on attribute sets within the HIR.
trait Attrs {
    /// Inserts a value into the attribute set at a given path.
    ///
    /// # Example
    /// `insert([a, b], value)` corresponds to `a.b = value;`.
    fn insert(
        &mut self,
        path: Vec<Attr>,
        value: ExprId,
        ctx: &mut impl DowngradeContext,
    ) -> Result<()>;

    /// Internal helper for recursively inserting an attribute.
    fn _insert(
        &mut self,
        path: impl Iterator<Item = Attr>,
        name: Attr,
        value: ExprId,
        ctx: &mut impl DowngradeContext,
    ) -> Result<()>;
}

impl Attrs for AttrSet {
    fn _insert(
        &mut self,
        mut path: impl Iterator<Item = Attr>,
        name: Attr,
        value: ExprId,
        ctx: &mut impl DowngradeContext,
    ) -> Result<()> {
        if let Some(attr) = path.next() {
            // If the path is not yet exhausted, we need to recurse deeper.
            match attr {
                Attr::Str(ident) => {
                    // If the next attribute is a static string.
                    if let Some(id) = self.stcs.get(&ident) {
                        // If a sub-attrset already exists, recurse into it.
                        ctx.with_expr_mut(id, |expr, ctx| {
                            expr.as_mut()
                                .try_unwrap_attr_set()
                                .map_err(|_| {
                                    // This path segment exists but is not an attrset.
                                    Error::DowngradeError(format!(
                                        r#"attribute '{}' already defined"#,
                                        format_symbol(&ident)
                                    ))
                                })
                                .and_then(|attrs| attrs._insert(path, name, value, ctx))
                        })?;
                    } else {
                        // Create a new sub-attrset because this path doesn't exist yet.
                        let mut attrs = AttrSet::default();
                        attrs._insert(path, name, value, ctx)?;
                        let attrs = ctx.new_expr(attrs.to_hir());
                        self.stcs.insert(ident, attrs);
                    }
                    Ok(())
                }
                Attr::Dynamic(dynamic) => {
                    // If the next attribute is a dynamic expression.
                    let mut attrs = AttrSet::default();
                    attrs._insert(path, name, value, ctx)?;
                    self.dyns.push((dynamic, ctx.new_expr(attrs.to_hir())));
                    Ok(())
                }
            }
        } else {
            // This is the final attribute in the path, so insert the value here.
            match name {
                Attr::Str(ident) => {
                    if self.stcs.insert(ident.clone(), value).is_some() {
                        return Err(Error::DowngradeError(format!(
                            r#"attribute '{}' already defined"#,
                            format_symbol(&ident)
                        )));
                    }
                }
                Attr::Dynamic(dynamic) => {
                    self.dyns.push((dynamic, value));
                }
            }
            Ok(())
        }
    }

    fn insert(
        &mut self,
        path: Vec<Attr>,
        value: ExprId,
        ctx: &mut impl DowngradeContext,
    ) -> Result<()> {
        let mut path = path.into_iter();
        // The last part of the path is the name of the attribute to be inserted.
        let name = path.next_back().unwrap();
        self._insert(path, name, value, ctx)
    }
}

#[derive(Debug)]
enum Param {
    /// A simple parameter, e.g., `x: ...`.
    Ident(String),
    /// A pattern-matching parameter (formals), e.g., `{ a, b ? 2, ... }@args: ...`.
    Formals {
        /// The individual formal parameters, with optional default values.
        formals: Vec<(String, Option<ExprId>)>,
        /// Whether an ellipsis (`...`) is present, allowing extra arguments.
        ellipsis: bool,
        /// An optional alias for the entire argument set, e.g., `args @ { ... }`.
        alias: Option<String>,
    },
}