use std::collections::BTreeMap;

use crate::mir::basic_block::BasicBlockId;
use crate::mir::builder::control_flow::edgecfg::api::exit_kind::ExitKind;
use crate::mir::builder::control_flow::edgecfg::api::frag::Frag;

/// Phase 281 P3: cleanup() Normal + Return exit wiring implementation
///
/// Wires cleanup Normal/Return exits to specified targets or propagates them upward.
///
/// # Contract (P3 Implementation)
///
/// **Input**:
/// - `main`: Main control flow (loop structure Frag)
/// - `cleanup_frag`: Exit handler (Normal/Return exits only, no wires/branches)
/// - `normal_target`: Where to wire Normal exits
///   - `Some(bb)`: Wire Normal → bb (internal closure, target = Some)
///   - `None`: Propagate Normal → wires (upward propagation, target = None)
/// - `ret_target`: Where to wire Return exits
///   - `Some(bb)`: Wire Return → bb (internal closure, target = Some)
///   - `None`: Propagate Return → wires (upward propagation, target = None)
///
/// **Output**:
/// - Frag with main's structure + cleanup's exits wired/propagated
///
/// **Invariants**:
/// - 1 block = 1 terminator (no duplicate BranchStubs)
/// - cleanup_frag must have empty wires/branches (Fail-Fast if not)
/// - cleanup_frag.exits must contain only Normal/Return (Fail-Fast for other kinds)
/// - normal_target=Some: Normal exits → wires (internal)
/// - normal_target=None: Normal exits → wires (target=None, propagate upward)
/// - ret_target=Some: Return exits → wires (internal)
/// - ret_target=None: Return exits → wires (target=None, propagate upward)
///
/// # Implementation Status
///
/// P3: Normal + Return wiring logic implemented
/// Future: Break/Continue/Unwind support (P4+)
///
/// # Migration Notes (Phase 264 → Phase 281)
///
/// Old signature (Phase 264): `cleanup(body: Frag, cleanup_block: BasicBlockId) -> Frag`
/// Phase 281 P1: `cleanup(main: Frag, cleanup: Frag) -> Result<Frag, String>`
/// Phase 281 P2: `cleanup(main: Frag, cleanup_frag: Frag, ret_target: Option<BasicBlockId>) -> Result<Frag, String>`
/// Phase 281 P3: `cleanup(main: Frag, cleanup_frag: Frag, normal_target: Option<BasicBlockId>, ret_target: Option<BasicBlockId>) -> Result<Frag, String>`
///
/// Rationale: Pattern6/7 require flexible exit wiring for Normal/Return exits.
/// cleanup_frag must be "exit-only" to prevent terminator confusion.
pub(crate) fn cleanup(
    main: Frag,
    cleanup_frag: Frag,
    normal_target: Option<BasicBlockId>,
    ret_target: Option<BasicBlockId>,
) -> Result<Frag, String> {
    // Phase 281 P3: Normal + Return exit wiring implementation
    // - Supported: Normal, Return exits
    // - Unsupported: Break, Continue, Unwind (Fail-Fast)

    let mut exits = BTreeMap::new();
    let mut wires = Vec::new();
    let mut branches = Vec::new();

    // Validate cleanup_frag structure (only exits allowed, no wires/branches)
    if !cleanup_frag.wires.is_empty() || !cleanup_frag.branches.is_empty() {
        return Err(format!(
            "compose::cleanup() Phase 281 P3: cleanup_frag must have empty wires/branches (only exits allowed), found {} wires, {} branches",
            cleanup_frag.wires.len(),
            cleanup_frag.branches.len()
        ));
    }

    // Validate cleanup_frag exits (only Normal + Return allowed in P3)
    for (kind, _) in &cleanup_frag.exits {
        match kind {
            ExitKind::Normal | ExitKind::Return => {} // OK
            _ => {
                return Err(format!(
                    "compose::cleanup() Phase 281 P3: unsupported exit kind {:?} in cleanup_frag (only Normal/Return allowed)",
                    kind
                ));
            }
        }
    }

    // Process cleanup Normal exits
    if let Some(normal_stubs) = cleanup_frag.exits.get(&ExitKind::Normal) {
        for mut stub in normal_stubs.clone() {
            match normal_target {
                Some(target_bb) => {
                    // Wire: Normal → target_bb (internal closure)
                    stub.target = Some(target_bb);
                    wires.push(stub);
                }
                None => {
                    // Propagate: Normal → wires (target=None, upward propagation)
                    stub.target = None;
                    wires.push(stub);
                }
            }
        }
    }

    // Process cleanup Return exits
    if let Some(return_stubs) = cleanup_frag.exits.get(&ExitKind::Return) {
        for mut stub in return_stubs.clone() {
            match ret_target {
                Some(target_bb) => {
                    // Wire: Return → target_bb (internal closure)
                    stub.target = Some(target_bb);
                    wires.push(stub);
                }
                None => {
                    // Propagate: Return → wires (target=None, will be emitted as Return terminator)
                    // Note: Return exits can have target=None in wires (Phase 267 special case)
                    stub.target = None;
                    wires.push(stub);
                }
            }
        }
    }

    // Preserve main's exits/wires/branches
    for (kind, stubs) in main.exits {
        exits.entry(kind).or_insert_with(Vec::new).extend(stubs);
    }
    wires.extend(main.wires);
    branches.extend(main.branches);

    Ok(Frag {
        entry: main.entry,  // Entry = main entry (header_bb)
        exits,
        wires,
        branches,
    })
}
