// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

//! Infers which variables belong to reactive scopes.
//!
//! Ported from TypeScript `src/ReactiveScopes/InferReactiveScopeVariables.ts`.
//!
//! This is the 1st of 4 passes that determine how to break a function into
//! discrete reactive scopes (independently memoizable units of code):
//! 1. InferReactiveScopeVariables (this pass, on HIR) determines operands that
//!    mutate together and assigns them a unique reactive scope.
//! 2. AlignReactiveScopesToBlockScopes aligns reactive scopes to block scopes.
//! 3. MergeOverlappingReactiveScopes ensures scopes do not overlap.
//! 4. BuildReactiveBlocks groups the statements for each scope.

use oxc_diagnostics::OxcDiagnostic;
use rustc_hash::FxHashMap;

use crate::diagnostics::ErrorCategory;
use crate::react_compiler_hir::environment::Environment;
use crate::react_compiler_hir::visitors;
use crate::react_compiler_hir::{
    DeclarationId, EvaluationOrder, HirFunction, IdentifierId, InstructionValue, Pattern, ScopeId,
    is_primitive_type,
};
use crate::react_compiler_utils::DisjointSet;
use oxc_span::Span;

// =============================================================================
// Public API
// =============================================================================

/// Infer reactive scope variables for a function.
///
/// For each mutable variable, infers a reactive scope which will construct that
/// variable. Variables that co-mutate are assigned to the same reactive scope.
///
/// Corresponds to TS `inferReactiveScopeVariables(fn: HIRFunction): void`.
pub fn infer_reactive_scope_variables(
    func: &mut HirFunction,
    env: &mut Environment,
) -> Result<(), OxcDiagnostic> {
    // Phase 1: find disjoint sets of co-mutating identifiers
    let mut scope_identifiers = find_disjoint_mutable_values(func, env);

    // Phase 2: assign scopes
    // Maps each group root identifier to the ScopeId assigned to that group.
    let mut scopes: FxHashMap<IdentifierId, ScopeState> = FxHashMap::default();

    scope_identifiers.for_each(|identifier_id, group_id| {
        let ident_range = env.identifiers[identifier_id].mutable_range.clone();
        let ident_span = env.identifiers[identifier_id].span;

        let state = scopes.entry(group_id).or_insert_with(|| {
            let scope_id = env.next_scope_id();
            // Initialize scope range from the first member
            let scope = &mut env.scopes[scope_id];
            scope.range = ident_range.clone();
            ScopeState { scope_id, span: ident_span }
        });

        // Update scope range
        let scope = &mut env.scopes[state.scope_id];

        // If this is not the first identifier (scope was already created), merge ranges
        if scope.range.start != ident_range.start || scope.range.end != ident_range.end {
            if scope.range.start == EvaluationOrder::UNSET {
                scope.range.start = ident_range.start;
            } else if ident_range.start != EvaluationOrder::UNSET {
                scope.range.start = scope.range.start.min(ident_range.start);
            }
            scope.range.end = scope.range.end.max(ident_range.end);
        }

        // Merge location
        state.span = merge_location(state.span, ident_span);

        // Assign the scope to this identifier
        let scope_id = state.scope_id;
        env.identifiers[identifier_id].scope = Some(scope_id);
    });

    // Set span on each scope
    for state in scopes.values() {
        env.scopes[state.scope_id].span = state.span;
    }

    // Update each identifier's mutable_range to match its scope's range
    for (&_identifier_id, state) in &scopes {
        let scope_range = env.scopes[state.scope_id].range.clone();
        // Find all identifiers with this scope and update their mutable_range
        // We iterate through all identifiers and check their scope
        for ident in &mut env.identifiers {
            if ident.scope == Some(state.scope_id) {
                ident.mutable_range = scope_range.clone();
            }
        }
    }

    // Validate scope ranges
    let mut max_instruction = EvaluationOrder::UNSET;
    for (_block_id, block) in &func.body.blocks {
        for instr_id in &block.instructions {
            let instr = &func.instructions[instr_id.index()];
            max_instruction = max_instruction.max(instr.id);
        }
        max_instruction = max_instruction.max(block.terminal.evaluation_order());
    }

    for state in scopes.values() {
        let scope = &env.scopes[state.scope_id];
        if scope.range.start == EvaluationOrder::UNSET
            || scope.range.end == EvaluationOrder::UNSET
            || max_instruction == EvaluationOrder::UNSET
            || scope.range.end > max_instruction + 1
        {
            return Err(ErrorCategory::Invariant.diagnostic(format!(
                "Invalid mutable range for scope: Scope @{} has range [{}:{}] but the valid range is [1:{}]",
                scope.id.index(),
                scope.range.start.index(),
                scope.range.end.index(),
                max_instruction.index() + 1,
            )));
        }
    }

    Ok(())
}

struct ScopeState {
    scope_id: ScopeId,
    span: Option<Span>,
}

/// Merge two source locations, preferring non-None values.
/// Corresponds to TS `mergeLocation`.
fn merge_location(l: Option<Span>, r: Option<Span>) -> Option<Span> {
    match (l, r) {
        (None, r) => r,
        (l, None) => l,
        (Some(l), Some(r)) => Some(Span::new(l.start.min(r.start), l.end.max(r.end))),
    }
}

// =============================================================================
// is_mutable / in_range helpers
// =============================================================================

// =============================================================================
// may_allocate
// =============================================================================

/// Check if an instruction may allocate. Corresponds to TS `mayAllocate`.
fn may_allocate(value: &InstructionValue, lvalue_type_is_primitive: bool) -> bool {
    match value {
        InstructionValue::Destructure { lvalue, .. } => {
            visitors::does_pattern_contain_spread_element(&lvalue.pattern)
        }
        InstructionValue::PostfixUpdate { .. }
        | InstructionValue::PrefixUpdate { .. }
        | InstructionValue::Await { .. }
        | InstructionValue::DeclareLocal { .. }
        | InstructionValue::DeclareContext { .. }
        | InstructionValue::StoreLocal { .. }
        | InstructionValue::LoadGlobal { .. }
        | InstructionValue::MetaProperty { .. }
        | InstructionValue::TypeCastExpression { .. }
        | InstructionValue::LoadLocal { .. }
        | InstructionValue::LoadContext { .. }
        | InstructionValue::StoreContext { .. }
        | InstructionValue::PropertyDelete { .. }
        | InstructionValue::ComputedLoad { .. }
        | InstructionValue::ComputedDelete { .. }
        | InstructionValue::JSXText { .. }
        | InstructionValue::TemplateLiteral { .. }
        | InstructionValue::Primitive { .. }
        | InstructionValue::GetIterator { .. }
        | InstructionValue::IteratorNext { .. }
        | InstructionValue::NextPropertyOf { .. }
        | InstructionValue::Debugger { .. }
        | InstructionValue::StartMemoize { .. }
        | InstructionValue::FinishMemoize { .. }
        | InstructionValue::UnaryExpression { .. }
        | InstructionValue::BinaryExpression { .. }
        | InstructionValue::PropertyLoad { .. }
        | InstructionValue::StoreGlobal { .. } => false,

        InstructionValue::TaggedTemplateExpression { .. }
        | InstructionValue::CallExpression { .. }
        | InstructionValue::MethodCall { .. } => !lvalue_type_is_primitive,

        InstructionValue::RegExpLiteral { .. }
        | InstructionValue::PropertyStore { .. }
        | InstructionValue::ComputedStore { .. }
        | InstructionValue::ArrayExpression { .. }
        | InstructionValue::JsxExpression { .. }
        | InstructionValue::JsxFragment { .. }
        | InstructionValue::NewExpression { .. }
        | InstructionValue::ObjectExpression { .. }
        | InstructionValue::ObjectMethod { .. }
        | InstructionValue::FunctionExpression { .. } => true,
    }
}

// =============================================================================
// Pattern helpers
// =============================================================================

/// Collect all Place identifiers from a destructure pattern.
/// Corresponds to TS `eachPatternOperand`.
fn each_pattern_operand(pattern: &Pattern) -> Vec<IdentifierId> {
    visitors::each_pattern_operand(pattern).into_iter().map(|p| p.identifier).collect()
}

/// Collect all operand identifiers from an instruction value.
/// Corresponds to TS `eachInstructionValueOperand`.
fn each_instruction_value_operand(
    value: &InstructionValue,
    env: &Environment,
) -> Vec<IdentifierId> {
    visitors::each_instruction_value_operand(value, env).into_iter().map(|p| p.identifier).collect()
}

// =============================================================================
// findDisjointMutableValues
// =============================================================================

/// Find disjoint sets of co-mutating identifier IDs.
///
/// Corresponds to TS `findDisjointMutableValues(fn: HIRFunction): DisjointSet<Identifier>`.
pub(crate) fn find_disjoint_mutable_values(
    func: &HirFunction,
    env: &Environment,
) -> DisjointSet<IdentifierId> {
    let mut scope_identifiers = DisjointSet::<IdentifierId>::new();
    let mut declarations: FxHashMap<DeclarationId, IdentifierId> = FxHashMap::default();

    for (_block_id, block) in &func.body.blocks {
        // Handle phi nodes
        for phi in &block.phis {
            let phi_id = phi.place.identifier;
            let phi_range = &env.identifiers[phi_id].mutable_range;
            let phi_decl_id = env.identifiers[phi_id].declaration_id;

            let first_instr_id = block
                .instructions
                .first()
                .map(|iid| func.instructions[iid.index()].id)
                .unwrap_or(block.terminal.evaluation_order());

            let is_phi_mutated_after_creation =
                phi_range.start + 1 != phi_range.end && phi_range.end > first_instr_id;
            // A phi operand defined at or after the phi's block is a loop
            // back-edge: the variable is reassigned within the loop (eg a
            // counter `a++` or `a = a + 1`). The reassignment must count as
            // the loop's scope reassigning the variable, so union the phi
            // with its operands and declaration. Otherwise the variable's
            // pre-loop value would become a dependency of the scope even
            // though the scope changes the value as it executes, making the
            // scope's dependencies unstable (the cached dependency would be
            // the post-loop value, which can never match the pre-loop value
            // compared at the top of the scope).
            let is_loop_carried_reassignment = !is_phi_mutated_after_creation
                && phi.operands.iter().any(|(_pred_id, operand)| {
                    env.identifiers[operand.identifier].mutable_range.start >= first_instr_id
                });
            if is_phi_mutated_after_creation || is_loop_carried_reassignment {
                let mut operands = vec![phi_id];
                if let Some(&decl_id) = declarations.get(&phi_decl_id) {
                    operands.push(decl_id);
                }
                for (_pred_id, phi_operand) in &phi.operands {
                    operands.push(phi_operand.identifier);
                }
                scope_identifiers.union(&operands);
            } else if env.config.enable_forest {
                for (_pred_id, phi_operand) in &phi.operands {
                    scope_identifiers.union(&[phi_id, phi_operand.identifier]);
                }
            }
        }

        // Handle instructions
        for instr_id in &block.instructions {
            let instr = &func.instructions[instr_id.index()];
            let mut operands: Vec<IdentifierId> = Vec::new();

            let lvalue_id = instr.lvalue.identifier;
            let lvalue_range = &env.identifiers[lvalue_id].mutable_range;
            let lvalue_type = &env.types[env.identifiers[lvalue_id].type_];
            let lvalue_type_is_primitive = is_primitive_type(lvalue_type);

            if lvalue_range.end > lvalue_range.start + 1
                || may_allocate(&instr.value, lvalue_type_is_primitive)
            {
                operands.push(lvalue_id);
            }

            match &instr.value {
                InstructionValue::DeclareLocal { lvalue, .. }
                | InstructionValue::DeclareContext { lvalue, .. } => {
                    let place_id = lvalue.place.identifier;
                    let decl_id = env.identifiers[place_id].declaration_id;
                    declarations.entry(decl_id).or_insert(place_id);
                }
                InstructionValue::StoreLocal { lvalue, value, .. }
                | InstructionValue::StoreContext { lvalue, value, .. } => {
                    let place_id = lvalue.place.identifier;
                    let decl_id = env.identifiers[place_id].declaration_id;
                    declarations.entry(decl_id).or_insert(place_id);

                    let place_range = &env.identifiers[place_id].mutable_range;
                    if place_range.end > place_range.start + 1 {
                        operands.push(place_id);
                    }

                    let value_range = &env.identifiers[value.identifier].mutable_range;
                    if value_range.contains(instr.id) && value_range.start > 0 {
                        operands.push(value.identifier);
                    }
                }
                InstructionValue::Destructure { lvalue, value, .. } => {
                    let pattern_places = each_pattern_operand(&lvalue.pattern);
                    for place_id in &pattern_places {
                        let decl_id = env.identifiers[*place_id].declaration_id;
                        declarations.entry(decl_id).or_insert(*place_id);

                        let place_range = &env.identifiers[*place_id].mutable_range;
                        if place_range.end > place_range.start + 1 {
                            operands.push(*place_id);
                        }
                    }

                    let value_range = &env.identifiers[value.identifier].mutable_range;
                    if value_range.contains(instr.id) && value_range.start > 0 {
                        operands.push(value.identifier);
                    }
                }
                InstructionValue::MethodCall { property, .. } => {
                    // For MethodCall: include all mutable operands plus the computed property
                    let all_operands = each_instruction_value_operand(&instr.value, env);
                    for op_id in &all_operands {
                        let op_range = &env.identifiers[*op_id].mutable_range;
                        if op_range.contains(instr.id) && op_range.start > 0 {
                            operands.push(*op_id);
                        }
                    }
                    // Ensure method property is in the same scope as the call
                    operands.push(property.identifier);
                }
                _ => {
                    // For all other instructions: include mutable operands
                    let all_operands = each_instruction_value_operand(&instr.value, env);
                    for op_id in &all_operands {
                        let op_range = &env.identifiers[*op_id].mutable_range;
                        if op_range.contains(instr.id) && op_range.start > 0 {
                            operands.push(*op_id);
                        }
                    }
                }
            }

            if !operands.is_empty() {
                scope_identifiers.union(&operands);
            }
        }
    }
    scope_identifiers
}
