#include "control_flow.h"
#include <algorithm>
#include <fmt/format.h>

namespace PSXRecomp {

ControlFlowAnalyzer::ControlFlowAnalyzer(const PS1Executable& exe) : exe_(exe) {}

bool ControlFlowAnalyzer::is_control_flow(uint32_t instr) {
    uint32_t opcode = (instr >> 26) & 0x3F;

    // SPECIAL opcode (0x00) - check funct for jr/jalr
    if (opcode == 0x00) {
        uint32_t funct = instr & 0x3F;
        return (funct == 0x08 || funct == 0x09); // jr or jalr
    }

    // REGIMM opcode (0x01) - bltz, bgez, bltzal, bgezal
    if (opcode == 0x01) {
        return true;
    }

    // Branch instructions
    if (opcode >= 0x04 && opcode <= 0x07) {
        return true; // beq, bne, blez, bgtz
    }

    // Branch likely variants (0x14-0x17)
    if (opcode >= 0x14 && opcode <= 0x17) {
        return true; // beql, bnel, blezl, bgtzl
    }

    // Jump instructions
    if (opcode == 0x02 || opcode == 0x03) {
        return true; // j, jal
    }

    return false;
}

uint32_t ControlFlowAnalyzer::get_branch_target(uint32_t pc, uint32_t instr) {
    // Branch target = (PC + 4) + (sign_extend(offset) << 2)
    int16_t offset = (int16_t)(instr & 0xFFFF);
    int32_t byte_offset = ((int32_t)offset) << 2;
    return (pc + 4) + byte_offset;
}

uint32_t ControlFlowAnalyzer::get_jump_target(uint32_t pc, uint32_t instr) {
    // Jump target = (PC+4)[31:28] | (instr_index << 2)
    uint32_t instr_index = instr & 0x03FFFFFF;
    uint32_t target = (instr_index << 2);
    uint32_t upper_bits = (pc + 4) & 0xF0000000;
    return upper_bits | target;
}

ControlFlowInstr ControlFlowAnalyzer::analyze_instruction(uint32_t addr, uint32_t instr) {
    ControlFlowInstr cf;
    cf.address = addr;
    cf.instruction = instr;
    cf.type = ControlFlowType::None;
    cf.target = 0;
    cf.has_delay_slot = false;
    cf.is_likely = false;

    uint32_t opcode = (instr >> 26) & 0x3F;
    uint32_t rs = (instr >> 21) & 0x1F;
    uint32_t rt = (instr >> 16) & 0x1F;
    uint32_t funct = instr & 0x3F;

    // SPECIAL opcode - jr/jalr
    if (opcode == 0x00) {
        if (funct == 0x08) { // jr
            cf.type = (rs == 31) ? ControlFlowType::Return : ControlFlowType::JumpRegister;
            cf.has_delay_slot = true;
            cf.mnemonic = "jr";
            return cf;
        } else if (funct == 0x09) { // jalr
            cf.type = ControlFlowType::JumpLinkReg;
            cf.has_delay_slot = true;
            cf.mnemonic = "jalr";
            return cf;
        }
    }

    // REGIMM opcode - bltz, bgez, bltzal, bgezal
    if (opcode == 0x01) {
        cf.type = ControlFlowType::Branch;
        cf.target = get_branch_target(addr, instr);
        cf.has_delay_slot = true;
        cf.mnemonic = (rt == 0x00) ? "bltz" :
                      (rt == 0x01) ? "bgez" :
                      (rt == 0x10) ? "bltzal" : "bgezal";
        return cf;
    }

    // Branch instructions
    switch (opcode) {
        case 0x04: // beq
            cf.type = ControlFlowType::Branch;
            cf.target = get_branch_target(addr, instr);
            cf.has_delay_slot = true;
            cf.mnemonic = "beq";
            break;

        case 0x05: // bne
            cf.type = ControlFlowType::Branch;
            cf.target = get_branch_target(addr, instr);
            cf.has_delay_slot = true;
            cf.mnemonic = "bne";
            break;

        case 0x06: // blez
            cf.type = ControlFlowType::Branch;
            cf.target = get_branch_target(addr, instr);
            cf.has_delay_slot = true;
            cf.mnemonic = "blez";
            break;

        case 0x07: // bgtz
            cf.type = ControlFlowType::Branch;
            cf.target = get_branch_target(addr, instr);
            cf.has_delay_slot = true;
            cf.mnemonic = "bgtz";
            break;

        // Branch likely variants
        case 0x14: // beql
            cf.type = ControlFlowType::Branch;
            cf.target = get_branch_target(addr, instr);
            cf.has_delay_slot = true;
            cf.is_likely = true;
            cf.mnemonic = "beql";
            break;

        case 0x15: // bnel
            cf.type = ControlFlowType::Branch;
            cf.target = get_branch_target(addr, instr);
            cf.has_delay_slot = true;
            cf.is_likely = true;
            cf.mnemonic = "bnel";
            break;

        case 0x16: // blezl
            cf.type = ControlFlowType::Branch;
            cf.target = get_branch_target(addr, instr);
            cf.has_delay_slot = true;
            cf.is_likely = true;
            cf.mnemonic = "blezl";
            break;

        case 0x17: // bgtzl
            cf.type = ControlFlowType::Branch;
            cf.target = get_branch_target(addr, instr);
            cf.has_delay_slot = true;
            cf.is_likely = true;
            cf.mnemonic = "bgtzl";
            break;

        // Jump instructions
        case 0x02: // j
            cf.type = ControlFlowType::Jump;
            cf.target = get_jump_target(addr, instr);
            cf.has_delay_slot = true;
            cf.mnemonic = "j";
            break;

        case 0x03: // jal
            cf.type = ControlFlowType::JumpLink;
            cf.target = get_jump_target(addr, instr);
            cf.has_delay_slot = true;
            cf.mnemonic = "jal";
            break;
    }

    return cf;
}

uint32_t ControlFlowAnalyzer::analysis_walk_hi(const Function& func) const {
    const uint32_t analysis_hi = exe_.analysis_end_address();
    const uint32_t walk_lo =
        func.alias_walk_lo ? func.alias_walk_lo : func.start_addr;
    // Only clamp when the clamp still leaves a non-empty walk. A function that
    // begins at or past the analysis bound is a degenerate discovery artefact;
    // leave its range alone so the emitter's mandatory-delay-slot check keeps
    // failing closed on it rather than silently producing an empty CFG.
    if (func.end_addr > analysis_hi && analysis_hi > walk_lo) return analysis_hi;
    return func.end_addr;
}

// NOTE: find_block_boundaries / build_basic_blocks / link_basic_blocks /
// detect_loops below are NOT on the live path — analyze_function does its own
// single-pass walk (see the comment there). They are kept in sync with the
// analysis bound anyway so wiring them up can never reintroduce the
// guard-word-as-block-leader defect.
std::set<uint32_t> ControlFlowAnalyzer::find_block_boundaries(const Function& func) {
    std::set<uint32_t> boundaries;
    const uint32_t walk_hi = analysis_walk_hi(func);

    // Function start is always a boundary
    boundaries.insert(func.start_addr);

    // Scan all instructions in function
    uint32_t addr = func.start_addr;
    while (addr < walk_hi) {
        auto instr_opt = exe_.read_word(addr);
        if (!instr_opt.has_value()) {
            break;
        }

        uint32_t instr = *instr_opt;

        if (is_control_flow(instr)) {
            ControlFlowInstr cf = analyze_instruction(addr, instr);

            // If branch/jump has a target, target is a boundary
            if (cf.target != 0 && cf.target >= func.start_addr && cf.target < walk_hi) {
                boundaries.insert(cf.target);
            }

            // Instruction after delay slot is a boundary (for branches, not jumps)
            if (cf.type == ControlFlowType::Branch ||
                cf.type == ControlFlowType::JumpLink ||
                cf.type == ControlFlowType::JumpLinkReg) {
                uint32_t fall_through = addr + 8; // After delay slot
                if (fall_through < walk_hi) {
                    boundaries.insert(fall_through);
                }
            }

            // For unconditional jump/return, next instruction is unreachable
            // (unless it's a jump target from elsewhere)
            if (cf.type == ControlFlowType::Jump ||
                cf.type == ControlFlowType::JumpRegister ||
                cf.type == ControlFlowType::Return) {
                uint32_t after_delay = addr + 8;
                if (after_delay < walk_hi) {
                    boundaries.insert(after_delay);
                }
            }
        }

        addr += 4;
    }

    return boundaries;
}

std::map<uint32_t, BasicBlock> ControlFlowAnalyzer::build_basic_blocks(
    const Function& func,
    const std::set<uint32_t>& boundaries) {

    std::map<uint32_t, BasicBlock> blocks;
    const uint32_t walk_hi = analysis_walk_hi(func);
    std::vector<uint32_t> boundary_vec(boundaries.begin(), boundaries.end());
    std::sort(boundary_vec.begin(), boundary_vec.end());

    for (size_t i = 0; i < boundary_vec.size(); i++) {
        BasicBlock block;
        block.start_addr = boundary_vec[i];

        if (i + 1 < boundary_vec.size()) {
            block.end_addr = boundary_vec[i + 1] - 4;
        } else {
            block.end_addr = walk_hi - 4;
        }

        if (block.end_addr < block.start_addr) continue;

        block.instruction_count = (block.end_addr - block.start_addr) / 4 + 1;
        block.is_entry = (block.start_addr == func.start_addr);
        block.is_exit = false;
        block.is_loop_header = false;

        // Analyze exit instruction (last instruction in block)
        auto exit_instr_opt = exe_.read_word(block.end_addr);
        if (exit_instr_opt.has_value()) {
            block.exit_instr = analyze_instruction(block.end_addr, *exit_instr_opt);

            // If last instruction is not control flow, check if it's a delay slot
            // (i.e., the instruction before it is a branch/jump with delay slot)
            if (block.exit_instr.type == ControlFlowType::None &&
                block.end_addr > block.start_addr) {
                auto branch_instr_opt = exe_.read_word(block.end_addr - 4);
                if (branch_instr_opt.has_value()) {
                    ControlFlowInstr branch_cf = analyze_instruction(block.end_addr - 4, *branch_instr_opt);
                    if (branch_cf.type != ControlFlowType::None && branch_cf.has_delay_slot) {
                        block.exit_instr = branch_cf;
                    }
                }
            }

            block.is_exit = (block.exit_instr.type == ControlFlowType::Return);
        } else {
            block.exit_instr.type = ControlFlowType::None;
        }

        blocks[block.start_addr] = block;
    }

    return blocks;
}

void ControlFlowAnalyzer::link_basic_blocks(std::map<uint32_t, BasicBlock>& blocks) {
    for (auto& [addr, block] : blocks) {
        ControlFlowInstr& cf = block.exit_instr;

        // Add successors based on control flow type
        if (cf.type == ControlFlowType::Branch) {
            // Branch has two successors: target and fall-through
            if (cf.target != 0 && blocks.count(cf.target)) {
                block.successors.push_back(cf.target);
                blocks[cf.target].predecessors.push_back(addr);
            }

            // Fall-through (after delay slot)
            uint32_t fall_through = cf.address + 8;
            if (blocks.count(fall_through)) {
                block.successors.push_back(fall_through);
                blocks[fall_through].predecessors.push_back(addr);
            }
        } else if (cf.type == ControlFlowType::Jump) {
            // Unconditional jump - only target
            if (cf.target != 0 && blocks.count(cf.target)) {
                block.successors.push_back(cf.target);
                blocks[cf.target].predecessors.push_back(addr);
            }
        } else if (cf.type == ControlFlowType::JumpLink || cf.type == ControlFlowType::JumpLinkReg) {
            // Function call - continues after delay slot
            uint32_t fall_through = cf.address + 8;
            if (blocks.count(fall_through)) {
                block.successors.push_back(fall_through);
                blocks[fall_through].predecessors.push_back(addr);
            }
        } else if (cf.type == ControlFlowType::JumpRegister && cf.type != ControlFlowType::Return) {
            // Register jump (not return) - unknown target, no successors
        } else if (cf.type == ControlFlowType::Return) {
            // Return - no successors
        } else {
            // No control flow - fall through to next block
            uint32_t next_block = block.end_addr + 4;
            if (blocks.count(next_block)) {
                block.successors.push_back(next_block);
                blocks[next_block].predecessors.push_back(addr);
            }
        }
    }
}

std::vector<std::pair<uint32_t, uint32_t>> ControlFlowAnalyzer::detect_loops(
    const std::map<uint32_t, BasicBlock>& blocks) {

    std::vector<std::pair<uint32_t, uint32_t>> loops;

    for (const auto& [addr, block] : blocks) {
        // Check each successor
        for (uint32_t successor : block.successors) {
            // Back edge: successor address < current block address
            if (successor <= addr) {
                loops.push_back({successor, addr}); // (header, back_edge_source)

                // Mark loop header
                if (blocks.count(successor)) {
                    const_cast<BasicBlock&>(blocks.at(successor)).is_loop_header = true;
                }
            }
        }
    }

    return loops;
}

ControlFlowGraph ControlFlowAnalyzer::analyze_function(const Function& func) {
    ControlFlowGraph cfg;
    // Alias entries cover their HOST's full range so backward branches, shared
    // epilogues, and jr jump tables (whose setup may precede the alias entry)
    // all stay intra-function. The emitted alias enters via goto to its
    // start_addr block; walk_lo is the host's start.
    uint32_t walk_lo = func.alias_walk_lo ? func.alias_walk_lo : func.start_addr;
    // Discovery/block-construction bound. Normally exactly func.end_addr; on
    // an overlay capture that carries trailing delay-slot guard words it stops
    // short of them, so the guard word can never become a block leader or a
    // block's exit instruction (its own delay slot does not exist). It stays
    // READABLE: a transfer at walk_hi - 4 still emits its slot from the guard.
    const uint32_t walk_hi = analysis_walk_hi(func);
    cfg.function_start = walk_lo;
    cfg.function_end = walk_hi;
    cfg.producer_lo = func.producer_lo;
    cfg.producer_hi = func.producer_hi;

    // Simplified single-pass analysis: scan instructions to find block boundaries
    // without building full CFG. This avoids the crash in build_basic_blocks.
    std::vector<uint32_t> boundary_vec;
    boundary_vec.push_back(walk_lo);
    if (func.alias_walk_lo) {
        // Every alias entry point of this host must be a block leader so the
        // shared body's entry switch has labels to land on — and so sibling
        // aliases produce identical CFGs.
        boundary_vec.push_back(func.start_addr);
        for (uint32_t e : func.alias_group_entries) {
            if (e >= walk_lo && e < walk_hi) boundary_vec.push_back(e);
        }
    }
    auto add_boundary = [&](uint32_t target) {
        if (target >= walk_lo && target < walk_hi) {
            boundary_vec.push_back(target);
        }
    };

    // Scan instructions for branches/jumps to find block starts
    for (uint32_t addr = walk_lo; addr < walk_hi; addr += 4) {
        auto instr_opt = exe_.read_word(addr);
        if (!instr_opt) continue;
        uint32_t instr = *instr_opt;

        if (is_control_flow(instr)) {
            ControlFlowInstr cf = analyze_instruction(addr, instr);

            if (cf.target != 0 && cf.target >= walk_lo && cf.target < walk_hi) {
                add_boundary(cf.target);
            }

            // After delay slot is a new block
            uint32_t after_delay = addr + 8;
            if (after_delay < walk_hi) {
                add_boundary(after_delay);
            }
        }
    }
    // Build blocks from boundaries
    std::sort(boundary_vec.begin(), boundary_vec.end());
    boundary_vec.erase(std::unique(boundary_vec.begin(), boundary_vec.end()), boundary_vec.end());

    for (size_t i = 0; i < boundary_vec.size(); i++) {
        BasicBlock block;
        block.start_addr = boundary_vec[i];
        block.end_addr = (i + 1 < boundary_vec.size()) ? boundary_vec[i + 1] - 4 : walk_hi - 4;

        if (block.end_addr < block.start_addr) continue;

        block.instruction_count = (block.end_addr - block.start_addr) / 4 + 1;
        block.is_entry = (block.start_addr == walk_lo);
        block.is_exit = false;
        block.is_loop_header = false;

        // Find exit instruction: scan backward from end to find last branch/jump
        block.exit_instr.type = ControlFlowType::None;
        block.exit_instr.address = block.end_addr;
        block.exit_instr.instruction = 0;
        block.exit_instr.target = 0;
        block.exit_instr.has_delay_slot = false;
        block.exit_instr.is_likely = false;

        // Check if the instruction before the last (the branch, not delay slot) is CF
        if (block.end_addr > block.start_addr) {
            auto prev_opt = exe_.read_word(block.end_addr - 4);
            if (prev_opt && is_control_flow(*prev_opt)) {
                block.exit_instr = analyze_instruction(block.end_addr - 4, *prev_opt);
                block.is_exit = (block.exit_instr.type == ControlFlowType::Return);
            }
        }
        // Also check last instruction itself
        if (block.exit_instr.type == ControlFlowType::None) {
            auto last_opt = exe_.read_word(block.end_addr);
            if (last_opt && is_control_flow(*last_opt)) {
                block.exit_instr = analyze_instruction(block.end_addr, *last_opt);
                block.is_exit = (block.exit_instr.type == ControlFlowType::Return);
            }
        }

        // Add successors — only targets within this function's block set
        auto in_func = [&](uint32_t addr) {
            return addr >= walk_lo &&
                   addr < walk_hi &&
                   std::binary_search(boundary_vec.begin(), boundary_vec.end(), addr);
        };
        if (block.exit_instr.type == ControlFlowType::Branch) {
            if (block.exit_instr.target != 0 && in_func(block.exit_instr.target))
                block.successors.push_back(block.exit_instr.target);
            uint32_t ft = block.exit_instr.address + 8;
            if (in_func(ft)) block.successors.push_back(ft);
        } else if (block.exit_instr.type == ControlFlowType::Jump) {
            if (block.exit_instr.target != 0 && in_func(block.exit_instr.target))
                block.successors.push_back(block.exit_instr.target);
        } else if (block.exit_instr.type == ControlFlowType::JumpLink ||
                   block.exit_instr.type == ControlFlowType::JumpLinkReg) {
            uint32_t ft = block.exit_instr.address + 8;
            if (in_func(ft)) block.successors.push_back(ft);
        } else if (block.exit_instr.type != ControlFlowType::Return &&
                   block.exit_instr.type != ControlFlowType::JumpRegister) {
            uint32_t next = block.end_addr + 4;
            if (in_func(next)) block.successors.push_back(next);
        }

        cfg.blocks[block.start_addr] = block;
    }
    // Detect back-edges (loops)
    for (const auto& [addr, block] : cfg.blocks) {
        for (uint32_t succ : block.successors) {
            if (succ <= addr) {
                cfg.loops.push_back({succ, addr});
                if (cfg.blocks.count(succ)) {
                    cfg.blocks[succ].is_loop_header = true;
                }
            }
        }
    }
    cfg.loop_count = static_cast<int>(cfg.loops.size());

    for (const auto& [addr, block] : cfg.blocks) {
        cfg.block_order.push_back(addr);
    }
    std::sort(cfg.block_order.begin(), cfg.block_order.end());

    return cfg;
}

std::map<uint32_t, ControlFlowGraph> ControlFlowAnalyzer::analyze_all_functions(
    const std::vector<Function>& functions) {

    std::map<uint32_t, ControlFlowGraph> all_cfgs;

    fmt::print("\n=== Control Flow Analysis ===\n\n");
    fmt::print("Analyzing {} functions...\n", functions.size());

    int total_blocks = 0;
    int total_loops = 0;

    for (const Function& func : functions) {
        ControlFlowGraph cfg = analyze_function(func);
        total_blocks += static_cast<int>(cfg.blocks.size());
        total_loops += cfg.loop_count;
        all_cfgs[func.start_addr] = cfg;
    }

    fmt::print("✓ Created {} basic blocks\n", total_blocks);
    fmt::print("✓ Detected {} loops\n", total_loops);
    fmt::print("✓ Average {:.1f} blocks per function\n\n",
               static_cast<double>(total_blocks) / functions.size());

    return all_cfgs;
}

} // namespace PSXRecomp