// 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.

//! Aligns scopes of object method values to that of their enclosing object expressions.
//! To produce a well-formed JS program in Codegen, object methods and object expressions
//! must be in the same ReactiveBlock as object method definitions must be inlined.
//!
//! Ported from TypeScript `src/ReactiveScopes/AlignObjectMethodScopes.ts`.

use rustc_hash::{FxHashMap, FxHashSet};
use std::cmp;
use std::mem::replace;

use crate::react_compiler_hir::MutableRangeId;
use crate::react_compiler_hir::environment::Environment;
use crate::react_compiler_hir::{
    EvaluationOrder, HirFunction, IdentifierId, InstructionValue, ObjectPropertyOrSpread, ScopeId,
};
use crate::react_compiler_ssa::enter_ssa::placeholder_function;
use crate::react_compiler_utils::DisjointSet;

// =============================================================================
// findScopesToMerge
// =============================================================================

/// Identifies ObjectMethod lvalue identifiers and then finds ObjectExpression
/// instructions whose operands reference those methods. Returns a disjoint set
/// of scopes that must be merged.
fn find_scopes_to_merge(func: &HirFunction, env: &Environment) -> DisjointSet<ScopeId> {
    let mut object_method_decls: FxHashSet<IdentifierId> = FxHashSet::default();
    let mut merged_scopes = DisjointSet::<ScopeId>::new();

    for (_block_id, block) in &func.body.blocks {
        for &instr_id in &block.instructions {
            let instr = &func.instructions[instr_id.index()];
            match &instr.value {
                InstructionValue::ObjectMethod { .. } => {
                    object_method_decls.insert(instr.lvalue.identifier);
                }
                InstructionValue::ObjectExpression { properties, .. } => {
                    for prop_or_spread in properties {
                        let operand_place = match prop_or_spread {
                            ObjectPropertyOrSpread::Property(prop) => &prop.place,
                            ObjectPropertyOrSpread::Spread(spread) => &spread.place,
                        };
                        if object_method_decls.contains(&operand_place.identifier) {
                            let operand_scope = env.identifiers[operand_place.identifier].scope;
                            let lvalue_scope = env.identifiers[instr.lvalue.identifier].scope;

                            // TS: Diagnostics.invariant(operandScope != null && lvalueScope != null, ...)
                            let operand_sid = operand_scope.expect(
                                "Internal error: Expected all ObjectExpressions and ObjectMethods to have non-null scope.",
                            );
                            let lvalue_sid = lvalue_scope.expect(
                                "Internal error: Expected all ObjectExpressions and ObjectMethods to have non-null scope.",
                            );
                            merged_scopes.union(&[operand_sid, lvalue_sid]);
                        }
                    }
                }
                _ => {}
            }
        }
    }

    merged_scopes
}

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

/// Aligns object method scopes so that ObjectMethod values and their enclosing
/// ObjectExpression share the same scope.
///
/// Corresponds to TS `alignObjectMethodScopes(fn: HIRFunction): void`.
pub fn align_object_method_scopes(func: &mut HirFunction, env: &mut Environment) {
    // Handle inner functions first (TS recurses before processing the outer function)
    for (_block_id, block) in &func.body.blocks {
        for &instr_id in &block.instructions {
            let instr = &func.instructions[instr_id.index()];
            match &instr.value {
                InstructionValue::FunctionExpression { lowered_func, .. }
                | InstructionValue::ObjectMethod { lowered_func, .. } => {
                    let func_id = lowered_func.func;
                    let mut inner_func =
                        replace(&mut env.functions[func_id], placeholder_function());
                    align_object_method_scopes(&mut inner_func, env);
                    env.functions[func_id] = inner_func;
                }
                _ => {}
            }
        }
    }

    let mut merged_scopes = find_scopes_to_merge(func, env);

    // Step 1: Merge affected scopes to their canonical root.
    // Use a FxHashMap to accumulate min/max across all scopes mapping to the same root,
    // matching TS behavior where root.range is updated in-place during iteration.
    let mut range_updates: FxHashMap<ScopeId, (EvaluationOrder, EvaluationOrder)> =
        FxHashMap::default();

    merged_scopes.for_each(|scope_id, root_id| {
        if scope_id == root_id {
            return;
        }
        let scope_range = env.scopes[scope_id].range.clone();
        let root_range = env.scopes[root_id].range.clone();

        let entry =
            range_updates.entry(root_id).or_insert_with(|| (root_range.start, root_range.end));
        entry.0 = cmp::min(entry.0, scope_range.start);
        entry.1 = cmp::max(entry.1, scope_range.end);
    });

    // Save original scope range IDs before updating
    let original_range_ids: FxHashMap<ScopeId, MutableRangeId> = range_updates
        .keys()
        .map(|&root_id| {
            let range_id = env.scopes[root_id].range.id;
            (root_id, range_id)
        })
        .collect();

    for (&root_id, (new_start, new_end)) in &range_updates {
        env.scopes[root_id].range.start = *new_start;
        env.scopes[root_id].range.end = *new_end;
    }

    // Sync identifier mutable_ranges that shared the old scope range.
    // Uses MutableRangeId for exact identity matching instead of value comparison.
    for ident in &mut env.identifiers {
        if let Some(scope_id) = ident.scope
            && let Some(&orig_range_id) = original_range_ids.get(&scope_id)
            && ident.mutable_range.id == orig_range_id
        {
            let new_range = &env.scopes[scope_id].range;
            ident.mutable_range.start = new_range.start;
            ident.mutable_range.end = new_range.end;
        }
    }

    // Step 2: Repoint identifiers whose scopes were merged
    // Build a map from old scope -> root scope for quick lookup
    let mut scope_remap: FxHashMap<ScopeId, ScopeId> = FxHashMap::default();
    merged_scopes.for_each(|scope_id, root_id| {
        if scope_id != root_id {
            scope_remap.insert(scope_id, root_id);
        }
    });

    for (_block_id, block) in &func.body.blocks {
        for &instr_id in &block.instructions {
            let lvalue_id = func.instructions[instr_id.index()].lvalue.identifier;

            if let Some(current_scope) = env.identifiers[lvalue_id].scope
                && let Some(&root) = scope_remap.get(&current_scope)
            {
                env.identifiers[lvalue_id].scope = Some(root);
            }
        }
    }
}
