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

//! Port of OutlineJsx from TypeScript.
//!
//! Outlines JSX expressions in callbacks into separate component functions.
//! This pass is conditional on `env.config.enable_jsx_outlining` (defaults to false).

use crate::react_compiler_utils::FxIndexSet;
use oxc_index::IndexVec;
use oxc_str::{Ident, IdentHashSet, format_ident};
use rustc_hash::{FxHashMap, FxHashSet};

use crate::react_compiler_hir::Effect;
use crate::react_compiler_hir::environment::Environment;
use crate::react_compiler_hir::{
    BasicBlock, BlockId, BlockKind, EvaluationOrder, FunctionId, HIR, HirFunction, IdentifierId,
    IdentifierName, Instruction, InstructionId, InstructionKind, InstructionValue, JsxAttribute,
    JsxTag, LValuePattern, NonLocalBinding, ObjectPattern, ObjectProperty, ObjectPropertyKey,
    ObjectPropertyOrSpread, ObjectPropertyType, ParamPattern, Pattern, Place, ReactFunctionType,
    ReturnVariant, Terminal,
};
use crate::react_compiler_ssa::enter_ssa::placeholder_function;
use crate::react_compiler_utils::FxIndexMap;
use std::mem::replace;

/// Outline JSX expressions in inner functions into separate outlined components.
///
/// Ported from TS `outlineJSX` in `Optimization/OutlineJsx.ts`.
pub fn outline_jsx<'a>(func: &mut HirFunction<'a>, env: &mut Environment<'a>) {
    let mut outlined_fns: IndexVec<FunctionId, HirFunction<'a>> = IndexVec::new();
    outline_jsx_impl(func, env, &mut outlined_fns);

    for outlined_fn in outlined_fns {
        env.outline_function(outlined_fn, Some(ReactFunctionType::Component));
    }
}

/// Data about a JSX instruction for outlining
struct JsxInstrInfo {
    instr_idx: usize, // index into func.instructions
    lvalue_id: IdentifierId,
    eval_order: EvaluationOrder,
}

struct OutlinedJsxAttribute<'a> {
    original_name: Ident<'a>,
    new_name: Ident<'a>,
    place: Place,
}

struct OutlinedResult<'a> {
    instrs: Vec<Instruction<'a>>,
    func: HirFunction<'a>,
}

fn outline_jsx_impl<'a>(
    func: &mut HirFunction<'a>,
    env: &mut Environment<'a>,
    outlined_fns: &mut IndexVec<FunctionId, HirFunction<'a>>,
) {
    // Collect LoadGlobal instructions (tag -> instr)
    let mut globals: FxHashMap<IdentifierId, usize> = FxHashMap::default(); // id -> instr_idx

    // Process each block
    let block_ids: Vec<BlockId> = func.body.blocks.keys().copied().collect();
    for block_id in &block_ids {
        let block = &func.body.blocks[block_id];
        let instr_ids = block.instructions.clone();

        let mut rewrite_instr: FxHashMap<EvaluationOrder, Vec<Instruction<'a>>> =
            FxHashMap::default();
        let mut jsx_group: Vec<JsxInstrInfo> = Vec::new();
        let mut children_ids: FxHashSet<IdentifierId> = FxHashSet::default();

        // First pass: collect all instruction info without borrowing func mutably
        enum InstrAction {
            LoadGlobal {
                lvalue_id: IdentifierId,
                instr_idx: usize,
            },
            FunctionExpr {
                func_id: FunctionId,
            },
            JsxExpr {
                lvalue_id: IdentifierId,
                instr_idx: usize,
                eval_order: EvaluationOrder,
                child_ids: Vec<IdentifierId>,
            },
            Other,
        }

        let mut actions: Vec<InstrAction> = Vec::new();
        for i in (0..instr_ids.len()).rev() {
            let iid = instr_ids[i];
            let instr = &func.instructions[iid.index()];
            let lvalue_id = instr.lvalue.identifier;

            match &instr.value {
                InstructionValue::LoadGlobal { .. } => {
                    actions.push(InstrAction::LoadGlobal { lvalue_id, instr_idx: iid.index() });
                }
                InstructionValue::FunctionExpression { lowered_func, .. } => {
                    actions.push(InstrAction::FunctionExpr { func_id: lowered_func.func });
                }
                InstructionValue::JsxExpression { children, .. } => {
                    let child_ids = children
                        .as_ref()
                        .map(|kids| kids.iter().map(|c| c.identifier).collect())
                        .unwrap_or_default();
                    actions.push(InstrAction::JsxExpr {
                        lvalue_id,
                        instr_idx: iid.index(),
                        eval_order: instr.id,
                        child_ids,
                    });
                }
                _ => {
                    actions.push(InstrAction::Other);
                }
            }
        }

        // Second pass: process actions
        for action in actions {
            match action {
                InstrAction::LoadGlobal { lvalue_id, instr_idx } => {
                    globals.insert(lvalue_id, instr_idx);
                }
                InstrAction::FunctionExpr { func_id } => {
                    let mut inner_func =
                        replace(&mut env.functions[func_id], placeholder_function());
                    outline_jsx_impl(&mut inner_func, env, outlined_fns);
                    env.functions[func_id] = inner_func;
                }
                InstrAction::JsxExpr { lvalue_id, instr_idx, eval_order, child_ids } => {
                    if !children_ids.contains(&lvalue_id) {
                        process_and_outline_jsx(
                            func,
                            env,
                            &mut jsx_group,
                            &globals,
                            &mut rewrite_instr,
                            outlined_fns,
                        );
                        jsx_group.clear();
                        children_ids.clear();
                    }
                    jsx_group.push(JsxInstrInfo { instr_idx, lvalue_id, eval_order });
                    for child_id in child_ids {
                        children_ids.insert(child_id);
                    }
                }
                InstrAction::Other => {}
            }
        }
        // Process remaining JSX group after the loop
        process_and_outline_jsx(
            func,
            env,
            &mut jsx_group,
            &globals,
            &mut rewrite_instr,
            outlined_fns,
        );
        if !rewrite_instr.is_empty() {
            let block = func.body.blocks.get_mut(block_id).unwrap();
            let old_instr_ids = block.instructions.clone();
            let mut new_instr_ids = Vec::new();
            for &iid in &old_instr_ids {
                let eval_order = func.instructions[iid.index()].id;
                if let Some(replacement_instrs) = rewrite_instr.get(&eval_order) {
                    // Add replacement instructions to the instruction table and reference them
                    for new_instr in replacement_instrs {
                        let new_idx = func.instructions.len();
                        func.instructions.push(new_instr.clone());
                        new_instr_ids.push(InstructionId::from_usize(new_idx));
                    }
                } else {
                    new_instr_ids.push(iid);
                }
            }
            let block = func.body.blocks.get_mut(block_id).unwrap();
            block.instructions = new_instr_ids;

            // Run dead code elimination after rewriting
            super::dead_code_elimination(func, env);
        }
    }
}

fn process_and_outline_jsx<'a>(
    func: &mut HirFunction<'a>,
    env: &mut Environment<'a>,
    jsx_group: &mut [JsxInstrInfo],
    globals: &FxHashMap<IdentifierId, usize>,
    rewrite_instr: &mut FxHashMap<EvaluationOrder, Vec<Instruction<'a>>>,
    outlined_fns: &mut IndexVec<FunctionId, HirFunction<'a>>,
) {
    if jsx_group.len() <= 1 {
        return;
    }
    // Sort by eval order ascending (TS: sort by a.id - b.id)
    jsx_group.sort_by_key(|j| j.eval_order);

    let result = process_jsx_group(func, env, jsx_group, globals);
    if let Some(result) = result {
        outlined_fns.push(result.func);
        // Map from the LAST JSX instruction's eval order to the replacement instructions
        // In the TS code, `state.jsx.at(0)` is the first element pushed during reverse iteration,
        // which is the last JSX in forward block order (highest eval order).
        // After sorting by eval_order ascending, that's jsx_group.last().
        let last_eval_order = jsx_group.last().unwrap().eval_order;
        rewrite_instr.insert(last_eval_order, result.instrs);
    }
}

fn process_jsx_group<'a>(
    func: &HirFunction<'a>,
    env: &mut Environment<'a>,
    jsx_group: &[JsxInstrInfo],
    globals: &FxHashMap<IdentifierId, usize>,
) -> Option<OutlinedResult<'a>> {
    // Only outline in callbacks, not top-level components
    if func.fn_type == ReactFunctionType::Component {
        return None;
    }

    let props = collect_props(func, env, jsx_group)?;

    let outlined_tag = env.generate_globally_unique_identifier_name(None);
    let new_instrs = emit_outlined_jsx(func, env, jsx_group, &props, outlined_tag)?;
    let outlined_fn = emit_outlined_fn(func, env, jsx_group, &props, globals)?;

    // Set the outlined function's id
    let mut outlined_fn = outlined_fn;
    outlined_fn.id = Some(outlined_tag);

    Some(OutlinedResult { instrs: new_instrs, func: outlined_fn })
}

fn collect_props<'a>(
    func: &HirFunction<'a>,
    env: &mut Environment<'a>,
    jsx_group: &[JsxInstrInfo],
) -> Option<Vec<OutlinedJsxAttribute<'a>>> {
    let mut id_counter = 1u32;
    let mut seen: IdentHashSet<'a> = IdentHashSet::default();
    let mut attributes = Vec::new();
    let jsx_ids: FxHashSet<IdentifierId> = jsx_group.iter().map(|j| j.lvalue_id).collect();

    let allocator = env.allocator;
    let mut generate_name = |old_name: Ident<'a>| -> Ident<'a> {
        let mut new_name = old_name;
        while seen.contains(&new_name) {
            new_name = format_ident!(allocator, "{old_name}{id_counter}");
            id_counter += 1;
        }
        seen.insert(new_name);
        // TS: env.programContext.addNewReference(newName)
        // We don't have programContext in Rust, but this is needed for unique name tracking
        new_name
    };

    for info in jsx_group {
        let instr = &func.instructions[info.instr_idx];
        if let InstructionValue::JsxExpression { props, children, .. } = &instr.value {
            for attr in props {
                match attr {
                    JsxAttribute::SpreadAttribute { .. } => return None,
                    JsxAttribute::Attribute { name, place } => {
                        let new_name = generate_name(*name);
                        attributes.push(OutlinedJsxAttribute {
                            original_name: *name,
                            new_name,
                            place: place.clone(),
                        });
                    }
                }
            }

            if let Some(kids) = children {
                for child in kids {
                    if jsx_ids.contains(&child.identifier) {
                        continue;
                    }
                    // Promote the child's identifier to a named temporary
                    let child_id = child.identifier;
                    let decl_id = env.identifiers[child_id].declaration_id;
                    if env.identifiers[child_id].name.is_none() {
                        env.identifiers[child_id].name = Some(IdentifierName::Promoted(
                            format_ident!(allocator, "#t{}", decl_id.index()),
                        ));
                    }

                    let child_name = match env.identifiers[child_id].name {
                        Some(IdentifierName::Named(n)) => n,
                        Some(IdentifierName::Promoted(n)) => n,
                        None => format_ident!(allocator, "#t{}", decl_id.index()),
                    };
                    let new_name = generate_name(Ident::from("t"));
                    attributes.push(OutlinedJsxAttribute {
                        original_name: child_name,
                        new_name,
                        place: child.clone(),
                    });
                }
            }
        }
    }

    Some(attributes)
}

fn emit_outlined_jsx<'a>(
    func: &HirFunction<'a>,
    env: &mut Environment<'a>,
    jsx_group: &[JsxInstrInfo],
    outlined_props: &[OutlinedJsxAttribute<'a>],
    outlined_tag: Ident<'a>,
) -> Option<Vec<Instruction<'a>>> {
    let props: Vec<JsxAttribute> = outlined_props
        .iter()
        .map(|p| JsxAttribute::Attribute { name: p.new_name, place: p.place.clone() })
        .collect();

    // Create LoadGlobal for the outlined component
    let load_id = env.next_identifier_id();
    // Promote it as a JSX tag temporary
    let decl_id = env.identifiers[load_id].declaration_id;
    env.identifiers[load_id].name =
        Some(IdentifierName::Promoted(format_ident!(env.allocator, "#T{}", decl_id.index())));

    let load_place =
        Place { identifier: load_id, effect: Effect::Unknown, reactive: false, span: None };

    let load_jsx = Instruction {
        id: EvaluationOrder::UNSET,
        lvalue: load_place.clone(),
        value: InstructionValue::LoadGlobal {
            binding: NonLocalBinding::ModuleLocal { name: outlined_tag },
            span: None,
        },
        span: None,
        effects: None,
    };

    // Create the replacement JsxExpression using the last JSX instruction's lvalue
    let last_info = jsx_group.last().unwrap();
    let last_instr = &func.instructions[last_info.instr_idx];
    let jsx_expr = Instruction {
        id: EvaluationOrder::UNSET,
        lvalue: last_instr.lvalue.clone(),
        value: InstructionValue::JsxExpression {
            tag: JsxTag::Place(load_place),
            props,
            children: None,
            span: None,
            opening_span: None,
            closing_span: None,
        },
        span: None,
        effects: None,
    };

    Some(vec![load_jsx, jsx_expr])
}

fn emit_outlined_fn<'a>(
    func: &HirFunction<'a>,
    env: &mut Environment<'a>,
    jsx_group: &[JsxInstrInfo],
    old_props: &[OutlinedJsxAttribute<'a>],
    globals: &FxHashMap<IdentifierId, usize>,
) -> Option<HirFunction<'a>> {
    let old_to_new_props = create_old_to_new_props_mapping(env, old_props);

    // Create props parameter
    let props_obj_id = env.next_identifier_id();
    let decl_id = env.identifiers[props_obj_id].declaration_id;
    env.identifiers[props_obj_id].name =
        Some(IdentifierName::Promoted(format_ident!(env.allocator, "#t{}", decl_id.index())));
    let props_obj =
        Place { identifier: props_obj_id, effect: Effect::Unknown, reactive: false, span: None };

    // Create destructure instruction
    let destructure_instr = emit_destructure_props(env, &props_obj, &old_to_new_props);

    // Emit load globals for JSX tags
    let load_global_instrs = emit_load_globals(func, jsx_group, globals)?;

    // Emit updated JSX instructions
    let updated_jsx_instrs = emit_updated_jsx(func, jsx_group, &old_to_new_props);

    // Build instructions list
    let mut instructions = Vec::new();
    instructions.push(destructure_instr);
    instructions.extend(load_global_instrs);
    instructions.extend(updated_jsx_instrs);

    // Build instruction table and instruction IDs
    let mut instr_table = Vec::new();
    let mut instr_ids = Vec::new();
    for instr in instructions {
        let idx = instr_table.len();
        instr_table.push(instr);
        instr_ids.push(InstructionId::from_usize(idx));
    }

    // Return terminal uses the last instruction's lvalue
    let last_lvalue = instr_table.last().unwrap().lvalue.clone();

    // Create return place
    let returns_id = env.next_identifier_id();
    let returns_place =
        Place { identifier: returns_id, effect: Effect::Unknown, reactive: false, span: None };

    let block = BasicBlock {
        kind: BlockKind::Block,
        id: BlockId::ENTRY,
        instructions: instr_ids,
        preds: FxIndexSet::default(),
        terminal: Terminal::Return {
            value: last_lvalue,
            return_variant: ReturnVariant::Explicit,
            id: EvaluationOrder::UNSET,
            span: None,
            effects: None,
        },
        phis: Vec::new(),
    };

    let mut blocks = FxIndexMap::default();
    blocks.insert(BlockId::ENTRY, block);

    let outlined_fn = HirFunction {
        id: None,
        name_hint: None,
        fn_type: ReactFunctionType::Other,
        params: vec![ParamPattern::Place(props_obj)],
        returns: returns_place,
        context: Vec::new(),
        body: HIR { entry: BlockId::ENTRY, blocks },
        instructions: instr_table,
        generator: false,
        is_async: false,
        directives: Vec::new(),
        aliasing_effects: Some(vec![]),
        span: None,
    };

    Some(outlined_fn)
}

fn emit_load_globals<'a>(
    func: &HirFunction<'a>,
    jsx_group: &[JsxInstrInfo],
    globals: &FxHashMap<IdentifierId, usize>,
) -> Option<Vec<Instruction<'a>>> {
    let mut instructions = Vec::new();
    for info in jsx_group {
        let instr = &func.instructions[info.instr_idx];
        if let InstructionValue::JsxExpression { tag: JsxTag::Place(tag_place), .. } = &instr.value
        {
            let global_instr_idx = globals.get(&tag_place.identifier)?;
            instructions.push(func.instructions[*global_instr_idx].clone());
        }
    }
    Some(instructions)
}

fn emit_updated_jsx<'a>(
    func: &HirFunction<'a>,
    jsx_group: &[JsxInstrInfo],
    old_to_new_props: &FxIndexMap<IdentifierId, OutlinedJsxAttribute<'a>>,
) -> Vec<Instruction<'a>> {
    let jsx_ids: FxHashSet<IdentifierId> = jsx_group.iter().map(|j| j.lvalue_id).collect();
    let mut new_instrs = Vec::new();

    for info in jsx_group {
        let instr = &func.instructions[info.instr_idx];
        if let InstructionValue::JsxExpression {
            tag,
            props,
            children,
            span,
            opening_span,
            closing_span,
        } = &instr.value
        {
            let mut new_props = Vec::new();
            for prop in props {
                // TS: invariant(prop.kind === 'JsxAttribute', ...)
                // Spread attributes would have caused collectProps to return null earlier
                let (name, place) = match prop {
                    JsxAttribute::Attribute { name, place } => (name, place),
                    JsxAttribute::SpreadAttribute { .. } => {
                        unreachable!("Expected only JsxAttribute, not spread")
                    }
                };
                if name == "key" {
                    continue;
                }
                // TS: invariant(newProp !== undefined, ...)
                let new_prop = old_to_new_props
                    .get(&place.identifier)
                    .expect("Expected a new property for identifier");
                new_props.push(JsxAttribute::Attribute {
                    name: new_prop.original_name,
                    place: new_prop.place.clone(),
                });
            }

            let new_children = children.as_ref().map(|kids| {
                kids.iter()
                    .map(|child| {
                        if jsx_ids.contains(&child.identifier) {
                            child.clone()
                        } else {
                            // TS: invariant(newChild !== undefined, ...)
                            let new_prop = old_to_new_props
                                .get(&child.identifier)
                                .expect("Expected a new prop for child identifier");
                            new_prop.place.clone()
                        }
                    })
                    .collect()
            });

            new_instrs.push(Instruction {
                id: instr.id,
                lvalue: instr.lvalue.clone(),
                value: InstructionValue::JsxExpression {
                    tag: tag.clone(),
                    props: new_props,
                    children: new_children,
                    span: *span,
                    opening_span: *opening_span,
                    closing_span: *closing_span,
                },
                span: instr.span,
                effects: instr.effects.clone(),
            });
        }
    }

    new_instrs
}

fn create_old_to_new_props_mapping<'a>(
    env: &mut Environment<'a>,
    old_props: &[OutlinedJsxAttribute<'a>],
) -> FxIndexMap<IdentifierId, OutlinedJsxAttribute<'a>> {
    let mut old_to_new = FxIndexMap::default();

    for old_prop in old_props {
        if old_prop.original_name == "key" {
            continue;
        }

        let new_id = env.next_identifier_id();
        env.identifiers[new_id].name = Some(IdentifierName::Named(old_prop.new_name));

        let new_place =
            Place { identifier: new_id, effect: Effect::Unknown, reactive: false, span: None };

        old_to_new.insert(
            old_prop.place.identifier,
            OutlinedJsxAttribute {
                original_name: old_prop.original_name,
                new_name: old_prop.new_name,
                place: new_place,
            },
        );
    }

    old_to_new
}

fn emit_destructure_props<'a>(
    env: &mut Environment<'a>,
    props_obj: &Place,
    old_to_new_props: &FxIndexMap<IdentifierId, OutlinedJsxAttribute<'a>>,
) -> Instruction<'a> {
    let mut properties = Vec::new();
    for prop in old_to_new_props.values() {
        properties.push(ObjectPropertyOrSpread::Property(ObjectProperty {
            key: ObjectPropertyKey::String { name: prop.new_name },
            property_type: ObjectPropertyType::Property,
            place: prop.place.clone(),
        }));
    }

    let lvalue_id = env.next_identifier_id();
    let lvalue =
        Place { identifier: lvalue_id, effect: Effect::Unknown, reactive: false, span: None };

    Instruction {
        id: EvaluationOrder::UNSET,
        lvalue,
        value: InstructionValue::Destructure {
            lvalue: LValuePattern {
                pattern: Pattern::Object(ObjectPattern { properties }),
                kind: InstructionKind::Let,
            },
            value: props_obj.clone(),
            span: None,
        },
        span: None,
        effects: None,
    }
}
