// Assume no parse error
#![allow(clippy::unwrap_used)]

use std::sync::Arc;

use hashbrown::hash_map::Entry;
use hashbrown::{HashMap, HashSet};
use rnix::ast;
use rowan::ast::AstNode;

use crate::error::{Error, Result};
use crate::ir::{Attr, AttrSet, ConcatStrings, ExprId, Ir, Select, Str, SymId};
use crate::value::format_symbol;

use super::*;

pub fn maybe_thunk(mut expr: ast::Expr, ctx: &mut impl DowngradeContext) -> Result<ExprId> {
    use ast::Expr::*;
    let expr = loop {
        expr = match expr {
            Paren(paren) => paren.expr().unwrap(),
            Root(root) => root.expr().unwrap(),
            expr => break expr,
        }
    };
    match expr {
        Error(error) => {
            let span = error.syntax().text_range();
            return Err(self::Error::downgrade_error(error.to_string())
                .with_span(span)
                .with_source(ctx.get_current_source().unwrap_or_else(|| Arc::from(""))));
        }
        Ident(ident) => return ident.downgrade(ctx),
        Literal(lit) => return lit.downgrade(ctx),
        Str(str) => return str.downgrade(ctx),
        Path(path) => return path.downgrade(ctx),

        _ => (),
    }
    let id = match expr {
        Apply(apply) => apply.downgrade(ctx),
        Assert(assert) => assert.downgrade(ctx),
        IfElse(ifelse) => ifelse.downgrade(ctx),
        Select(select) => select.downgrade(ctx),
        Lambda(lambda) => lambda.downgrade(ctx),
        LegacyLet(let_) => let_.downgrade(ctx),
        LetIn(letin) => letin.downgrade(ctx),
        List(list) => list.downgrade(ctx),
        BinOp(op) => op.downgrade(ctx),
        AttrSet(attrs) => attrs.downgrade(ctx),
        UnaryOp(op) => op.downgrade(ctx),
        With(with) => with.downgrade(ctx),
        HasAttr(has) => has.downgrade(ctx),

        _ => unreachable!(),
    }?;
    Ok(ctx.new_expr(
        Thunk {
            inner: id,
            // span: ctx.get_span(id),
            // FIXME: span
            span: synthetic_span()
        }
        .to_ir(),
    ))
}

/// Downgrades the entries of an attribute set.
/// This handles `inherit` and `attrpath = value;` entries.
pub fn downgrade_attrs(
    attrs: impl ast::HasEntry + AstNode,
    ctx: &mut impl DowngradeContext,
) -> Result<AttrSet> {
    let entries = attrs.entries();
    let mut attrs = AttrSet {
        stcs: HashMap::new(),
        dyns: Vec::new(),
        span: attrs.syntax().text_range(),
    };

    for entry in entries {
        match entry {
            ast::Entry::Inherit(inherit) => downgrade_inherit(inherit, &mut attrs.stcs, ctx)?,
            ast::Entry::AttrpathValue(value) => downgrade_attrpathvalue(value, &mut attrs, ctx)?,
        }
    }

    Ok(attrs)
}

/// Downgrades attribute set entries for a `let...in` expression.
/// This is a stricter version of `downgrade_attrs` that disallows dynamic attributes,
/// as `let` bindings must be statically known.
pub fn downgrade_static_attrs(
    attrs: impl ast::HasEntry + AstNode,
    ctx: &mut impl DowngradeContext,
) -> Result<HashMap<SymId, ExprId>> {
    let entries = attrs.entries();
    let mut attrs = AttrSet {
        stcs: HashMap::new(),
        dyns: Vec::new(),
        span: attrs.syntax().text_range(),
    };

    for entry in entries {
        match entry {
            ast::Entry::Inherit(inherit) => downgrade_inherit(inherit, &mut attrs.stcs, ctx)?,
            ast::Entry::AttrpathValue(value) => {
                downgrade_static_attrpathvalue(value, &mut attrs, ctx)?
            }
        }
    }

    Ok(attrs.stcs)
}

/// Downgrades an `inherit` statement.
/// `inherit (from) a b;` is translated into `a = from.a; b = from.b;`.
/// `inherit a b;` is translated into `a = a; b = b;` (i.e., bringing variables into scope).
pub fn downgrade_inherit(
    inherit: ast::Inherit,
    stcs: &mut HashMap<SymId, ExprId>,
    ctx: &mut impl DowngradeContext,
) -> Result<()> {
    // Downgrade the `from` expression if it exists.
    let from = if let Some(from) = inherit.from() {
        Some(from.expr().unwrap().downgrade(ctx)?)
    } else {
        None
    };
    for attr in inherit.attrs() {
        let span = attr.syntax().text_range();
        let ident = match downgrade_attr(attr, ctx)? {
            Attr::Str(ident) => ident,
            _ => {
                // `inherit` does not allow dynamic attributes.
                return Err(Error::downgrade_error(
                    "dynamic attributes not allowed in inherit".to_string(),
                )
                .with_span(span)
                .with_source(ctx.get_current_source().unwrap_or_else(|| Arc::from(""))));
            }
        };
        let expr = if let Some(expr) = from {
            let select_expr = ctx.new_expr(
                Select {
                    expr,
                    attrpath: vec![Attr::Str(ident)],
                    default: None,
                    span,
                }
                .to_ir(),
            );
            ctx.new_expr(
                Thunk {
                    inner: select_expr,
                    span,
                }
                .to_ir(),
            )
        } else {
            ctx.lookup(ident, span)?
        };
        match stcs.entry(ident) {
            Entry::Occupied(occupied) => {
                return Err(Error::downgrade_error(format!(
                    "attribute '{}' already defined",
                    format_symbol(ctx.get_sym(*occupied.key()))
                ))
                .with_span(span)
                .with_source(ctx.get_current_source().unwrap_or_else(|| Arc::from(""))));
            }
            Entry::Vacant(vacant) => vacant.insert(expr),
        };
    }
    Ok(())
}

/// Downgrades a single attribute key (part of an attribute path).
/// An attribute can be a static identifier, an interpolated string, or a dynamic expression.
pub fn downgrade_attr(attr: ast::Attr, ctx: &mut impl DowngradeContext) -> Result<Attr> {
    use ast::Attr::*;
    use ast::InterpolPart::*;
    let span = attr.syntax().text_range();
    match attr {
        Ident(ident) => Ok(Attr::Str(ctx.new_sym(ident.to_string()))),
        Str(string) => {
            let parts = string.normalized_parts();
            if parts.is_empty() {
                Ok(Attr::Str(ctx.new_sym("".to_string())))
            } else if parts.len() == 1 {
                // If the string has only one part, it's either a literal or a single interpolation.
                match parts.into_iter().next().unwrap() {
                    Literal(ident) => Ok(Attr::Str(ctx.new_sym(ident))),
                    Interpolation(interpol) => {
                        Ok(Attr::Dynamic(interpol.expr().unwrap().downgrade(ctx)?))
                    }
                }
            } else {
                // If the string has multiple parts, it's an interpolated string that must be concatenated.
                let parts = parts
                    .into_iter()
                    .map(|part| match part {
                        Literal(lit) => Ok(ctx.new_expr(self::Str { val: lit, span }.to_ir())),
                        Interpolation(interpol) => interpol.expr().unwrap().downgrade(ctx),
                    })
                    .collect::<Result<Vec<_>>>()?;
                Ok(Attr::Dynamic(
                    ctx.new_expr(ConcatStrings { parts, span }.to_ir()),
                ))
            }
        }
        Dynamic(dynamic) => Ok(Attr::Dynamic(dynamic.expr().unwrap().downgrade(ctx)?)),
    }
}

/// Downgrades an attribute path (e.g., `a.b."${c}".d`) into a `Vec<Attr>`.
pub fn downgrade_attrpath(
    attrpath: ast::Attrpath,
    ctx: &mut impl DowngradeContext,
) -> Result<Vec<Attr>> {
    attrpath
        .attrs()
        .map(|attr| downgrade_attr(attr, ctx))
        .collect::<Result<Vec<_>>>()
}

/// Downgrades an `attrpath = value;` expression and inserts it into an `AttrSet`.
pub fn downgrade_attrpathvalue(
    value: ast::AttrpathValue,
    attrs: &mut AttrSet,
    ctx: &mut impl DowngradeContext,
) -> Result<()> {
    let path = downgrade_attrpath(value.attrpath().unwrap(), ctx)?;
    let value = maybe_thunk(value.value().unwrap(), ctx)?;
    attrs.insert(path, value, ctx)
}

/// A stricter version of `downgrade_attrpathvalue` for `let...in` bindings.
/// It ensures that the attribute path contains no dynamic parts.
pub fn downgrade_static_attrpathvalue(
    value: ast::AttrpathValue,
    attrs: &mut AttrSet,
    ctx: &mut impl DowngradeContext,
) -> Result<()> {
    let attrpath_node = value.attrpath().unwrap();
    let path = downgrade_attrpath(attrpath_node.clone(), ctx)?;
    if path.iter().any(|attr| matches!(attr, Attr::Dynamic(_))) {
        return Err(Error::downgrade_error(
            "dynamic attributes not allowed in let bindings".to_string(),
        )
        .with_span(attrpath_node.syntax().text_range())
        .with_source(ctx.get_current_source().unwrap_or_else(|| Arc::from(""))));
    }
    let value = value.value().unwrap().downgrade(ctx)?;
    attrs.insert(path, value, ctx)
}

pub struct PatternBindings {
    pub body: ExprId,
    pub scc_info: SccInfo,
    pub required_params: Vec<SymId>,
    pub allowed_params: Option<HashSet<SymId>>,
}

/// Helper function for Lambda pattern parameters with SCC analysis.
/// Processes pattern entries like `{ a, b ? 2, ... }@alias` and creates optimized bindings.
///
/// # Parameters
/// - `pat_entries`: Iterator over pattern entries from the AST
/// - `alias`: Optional alias symbol (from @alias syntax)
/// - `arg`: The argument expression to extract from
///
/// Returns a tuple of (binding slots, body, SCC info, required params, allowed params)
pub fn downgrade_pattern_bindings<Ctx>(
    pat_entries: impl Iterator<Item = ast::PatEntry>,
    alias: Option<SymId>,
    arg: ExprId,
    has_ellipsis: bool,
    ctx: &mut Ctx,
    body_fn: impl FnOnce(&mut Ctx, &[SymId]) -> Result<ExprId>,
) -> Result<PatternBindings>
where
    Ctx: DowngradeContext,
{
    let mut param_syms = Vec::new();
    let mut param_defaults = Vec::new();
    let mut param_spans = Vec::new();
    let mut seen_params = HashSet::new();

    for entry in pat_entries {
        let sym = ctx.new_sym(entry.ident().unwrap().to_string());
        let span = entry.ident().unwrap().syntax().text_range();

        if !seen_params.insert(sym) {
            return Err(Error::downgrade_error(format!(
                "duplicate parameter '{}'",
                format_symbol(ctx.get_sym(sym))
            ))
            .with_span(span)
            .with_source(ctx.get_current_source().unwrap_or_else(|| Arc::from(""))));
        }

        let default_ast = entry.default();
        param_syms.push(sym);
        param_defaults.push(default_ast);
        param_spans.push(span);
    }

    let mut binding_keys: Vec<SymId> = param_syms.clone();
    if let Some(alias_sym) = alias {
        binding_keys.push(alias_sym);
    }

    let required: Vec<SymId> = param_syms
        .iter()
        .zip(param_defaults.iter())
        .filter_map(|(&sym, default)| if default.is_none() { Some(sym) } else { None })
        .collect();

    let allowed: Option<HashSet<SymId>> = if has_ellipsis {
        None
    } else {
        Some(param_syms.iter().copied().collect())
    };

    // Get the owner from outer tracker's current_binding
    let owner = ctx.get_current_binding();

    let (scc_info, body) = downgrade_bindings_generic_with_owner(
        ctx,
        binding_keys,
        |ctx, sym_to_slot| {
            let mut bindings = HashMap::new();

            for ((sym, default_ast), span) in param_syms
                .iter()
                .zip(param_defaults.iter())
                .zip(param_spans.iter())
            {
                let slot = *sym_to_slot.get(sym).unwrap();
                ctx.set_current_binding(Some(slot));

                let default = if let Some(default_expr) = default_ast {
                    Some(default_expr.clone().downgrade(ctx)?)
                } else {
                    None
                };

                let select_expr = ctx.new_expr(
                    Select {
                        expr: arg,
                        attrpath: vec![Attr::Str(*sym)],
                        default,
                        span: *span,
                    }
                    .to_ir(),
                );
                bindings.insert(*sym, select_expr);
                ctx.set_current_binding(None);
            }

            if let Some(alias_sym) = alias {
                bindings.insert(alias_sym, arg);
            }

            Ok(bindings)
        },
        body_fn,
        owner, // Pass the owner to track cross-scope dependencies
    )?;

    Ok(PatternBindings {
        body,
        scc_info,
        required_params: required,
        allowed_params: allowed,
    })
}

/// Generic helper function to downgrade bindings with SCC analysis.
/// This is the core logic for let bindings, extracted for reuse.
///
/// # Parameters
/// - `binding_keys`: The symbols for all bindings
/// - `compute_bindings_fn`: Called in let scope with sym_to_slot mapping to compute binding values
/// - `body_fn`: Called in let scope to compute the body expression
///
/// Returns a tuple of (binding slots, body result, SCC info)
pub fn downgrade_bindings_generic<Ctx, B, F>(
    ctx: &mut Ctx,
    binding_keys: Vec<SymId>,
    compute_bindings_fn: B,
    body_fn: F,
) -> Result<(SccInfo, ExprId)>
where
    Ctx: DowngradeContext,
    B: FnOnce(&mut Ctx, &HashMap<SymId, ExprId>) -> Result<HashMap<SymId, ExprId>>,
    F: FnOnce(&mut Ctx, &[SymId]) -> Result<ExprId>,
{
    downgrade_bindings_generic_with_owner(ctx, binding_keys, compute_bindings_fn, body_fn, None)
}

pub fn downgrade_bindings_generic_with_owner<Ctx, B, F>(
    ctx: &mut Ctx,
    binding_keys: Vec<SymId>,
    compute_bindings_fn: B,
    body_fn: F,
    owner: Option<ExprId>,
) -> Result<(SccInfo, ExprId)>
where
    Ctx: DowngradeContext,
    B: FnOnce(&mut Ctx, &HashMap<SymId, ExprId>) -> Result<HashMap<SymId, ExprId>>,
    F: FnOnce(&mut Ctx, &[SymId]) -> Result<ExprId>,
{
    let slots: Vec<_> = ctx.reserve_slots(binding_keys.len()).collect();
    let let_bindings: HashMap<_, _> = binding_keys
        .iter()
        .copied()
        .zip(slots.iter().copied())
        .collect();

    if let Some(owner_binding) = owner {
        ctx.push_dep_tracker_with_owner(&slots, owner_binding);
    } else {
        ctx.push_dep_tracker(&slots);
    }

    ctx.with_let_scope(let_bindings.clone(), |ctx| {
        let bindings = compute_bindings_fn(ctx, &let_bindings)?;

        let scc_info = ctx.pop_dep_tracker()?;

        for (sym, slot) in binding_keys.iter().copied().zip(slots.iter()) {
            if let Some(&expr) = bindings.get(&sym) {
                ctx.replace_expr(
                    *slot,
                    Thunk {
                        inner: expr,
                        // span: ctx.get_span(expr),
                        // FIXME: span
                        span: synthetic_span()
                    }
                    .to_ir(),
                );
            } else {
                return Err(Error::downgrade_error(format!(
                    "binding '{}' not found",
                    format_symbol(ctx.get_sym(sym))
                ))
                .with_span(synthetic_span())
                .with_source(ctx.get_current_source().unwrap_or_else(|| Arc::from(""))));
            }
        }

        let body = body_fn(ctx, &binding_keys)?;

        Ok((scc_info, body))
    })
}

/// Helper function to downgrade entries with let bindings semantics.
/// This extracts common logic for both `rec` attribute sets and `let...in` expressions.
///
/// Returns a tuple of (binding slots, body result, SCC info) where:
/// - binding slots: pre-allocated expression slots for the bindings
/// - body result: the result of calling `body_fn` in the let scope
/// - SCC info: strongly connected components information for optimization
pub fn downgrade_let_bindings<Ctx, F>(
    entries: Vec<ast::Entry>,
    ctx: &mut Ctx,
    body_fn: F,
) -> Result<(SccInfo, ExprId)>
where
    Ctx: DowngradeContext,
    F: FnOnce(&mut Ctx, &[SymId]) -> Result<ExprId>,
{
    let mut binding_syms = HashSet::new();

    for entry in &entries {
        match entry {
            ast::Entry::Inherit(inherit) => {
                for attr in inherit.attrs() {
                    if let ast::Attr::Ident(ident) = attr {
                        let sym = ctx.new_sym(ident.to_string());
                        if !binding_syms.insert(sym) {
                            return Err(Error::downgrade_error(format!(
                                "attribute '{}' already defined",
                                format_symbol(ctx.get_sym(sym))
                            ))
                            .with_span(ident.syntax().text_range())
                            .with_source(ctx.get_current_source().unwrap_or_else(|| Arc::from(""))));
                        }
                    }
                }
            }
            ast::Entry::AttrpathValue(value) => {
                let attrpath = value.attrpath().unwrap();
                let attrs_vec: Vec<_> = attrpath.attrs().collect();

                // Only check for duplicate definitions if this is a top-level binding (path length == 1)
                // For nested paths (e.g., types.a, types.b), they will be merged into the same attrset
                if attrs_vec.len() == 1 {
                    if let Some(ast::Attr::Ident(ident)) = attrs_vec.first() {
                        let sym = ctx.new_sym(ident.to_string());
                        if !binding_syms.insert(sym) {
                            return Err(Error::downgrade_error(format!(
                                "attribute '{}' already defined",
                                format_symbol(ctx.get_sym(sym))
                            ))
                            .with_span(ident.syntax().text_range())
                            .with_source(ctx.get_current_source().unwrap_or_else(|| Arc::from(""))));
                        }
                    }
                } else if attrs_vec.len() > 1 {
                    // For nested paths, just record the first-level name without checking duplicates
                    if let Some(ast::Attr::Ident(ident)) = attrs_vec.first() {
                        let sym = ctx.new_sym(ident.to_string());
                        binding_syms.insert(sym);
                    }
                }
            }
        }
    }

    let binding_keys: Vec<_> = binding_syms.into_iter().collect();

    downgrade_bindings_generic(
        ctx,
        binding_keys,
        |ctx, sym_to_slot| {
            let mut temp_attrs = AttrSet {
                stcs: HashMap::new(),
                dyns: Vec::new(),
                span: synthetic_span()
            };

            for entry in entries {
                match entry {
                    ast::Entry::Inherit(inherit) => {
                        for attr in inherit.attrs() {
                            if let ast::Attr::Ident(ident) = attr {
                                let sym = ctx.new_sym(ident.to_string());
                                let slot = *sym_to_slot.get(&sym).unwrap();
                                ctx.set_current_binding(Some(slot));
                            }
                        }
                        downgrade_inherit(inherit, &mut temp_attrs.stcs, ctx)?;
                        ctx.set_current_binding(None);
                    }
                    ast::Entry::AttrpathValue(value) => {
                        let attrpath = value.attrpath().unwrap();
                        if let Some(first_attr) = attrpath.attrs().next()
                            && let ast::Attr::Ident(ident) = first_attr
                        {
                            let sym = ctx.new_sym(ident.to_string());
                            let slot = *sym_to_slot.get(&sym).unwrap();
                            ctx.set_current_binding(Some(slot));
                        }
                        downgrade_static_attrpathvalue(value, &mut temp_attrs, ctx)?;
                        ctx.set_current_binding(None);
                    }
                }
            }

            Ok(temp_attrs.stcs)
        },
        body_fn,
    )
}