#include "cfg/builder.hpp"
#include <queue>
#include <algorithm>

static bool isMemOp(const string& op)
{
    return op.find('[') != string::npos || op.find('+') != string::npos;
}

static bool regMatch(const string& op, const string& target)
{
    auto norm = [](const string& s) -> string {
        if (s.size() >= 3 && s[0] == 'r' && isdigit(s[1])) return s.substr(1);
        if (s.size() >= 3 && s[0] == 'e' && isdigit(s[1])) return s.substr(1);
        return s;
    };
    return norm(op) == norm(target) || op == target;
}

CFGBuilder::CFGBuilder(Binary* binary, Disassembler* disasm)
    : m_bin(binary), m_dis(disasm) {}

bool CFGBuilder::isVisited(addr_t addr)
{
    return m_blocks.count(addr) > 0;
}

void CFGBuilder::markVisited(addr_t addr)
{
    m_blocks.insert(addr);
}

void CFGBuilder::buildGotMap()
{
    auto imports = m_bin->getImportedFunctions();
    for (const auto& imp : imports)
    {
        if (imp.address)
            m_got[imp.address] = {imp.name, imp.library};
    }
}

addr_t CFGBuilder::resolveGOT(const Instruction& inst)
{
    if (!m_dis->isIndirectBranch(inst)) return 0;
    int64_t disp = m_dis->getRIPDisp(inst);
    if (disp == 0) return 0;

    addr_t got = inst.address + inst.size + disp;
    auto it = m_got.find(got);
    if (it != m_got.end())
    {
        if (!it->second.first.empty()) return got;
    }

    auto bytes = m_bin->readBytes(got, 8);
    if (bytes.size() < 4) return 0;

    uint64_t target = 0;
    if (bytes.size() >= 8)
    {
        target = bytes[0] | ((uint64_t)bytes[1] << 8) | ((uint64_t)bytes[2] << 16) |
                 ((uint64_t)bytes[3] << 24) | ((uint64_t)bytes[4] << 32) |
                 ((uint64_t)bytes[5] << 40) | ((uint64_t)bytes[6] << 48) |
                 ((uint64_t)bytes[7] << 56);
    }
    else
    {
        target = bytes[0] | (bytes[1] << 8) | (bytes[2] << 16) | (bytes[3] << 24);
    }

    return target;
}

addr_t CFGBuilder::resolveJumpTable(const vector<Instruction>& block, const Instruction& inst)
{
    size_t n = block.size();
    if (n < 3) return 0;

    int idx = -1;
    for (int i = (int)n - 1; i >= 0; i--)
    {
        if (block[i].address == inst.address) { idx = i; break; }
    }
    if (idx < 2) return 0;

    for (int i = idx - 1; i >= max(0, idx - 8); i--)
    {
        if (block[i].mnemonic == "lea")
        {
            int64_t disp = m_dis->getRIPDisp(block[i]);
            if (disp == 0) continue;
            addr_t table = block[i].address + block[i].size + disp;

            int64_t bound = -1;
            for (int j = i - 1; j >= max(0, i - 4); j--)
            {
                if (block[j].mnemonic == "cmp")
                {
                    size_t pos = block[j].operands.find("0x");
                    if (pos != string::npos)
                    {
                        try { bound = stoll(block[j].operands.substr(pos), nullptr, 16); }
                        catch (...) {}
                    }
                    break;
                }
            }
            if (bound <= 0) bound = 32;

            for (int64_t e = 0; e < bound && e < 256; e++)
            {
                auto ent = m_bin->readBytes(table + e * 8, 8);
                if (ent.size() < 4) break;
                uint64_t entry = 0;
                if (ent.size() >= 8)
                    entry = ent[0] | ((uint64_t)ent[1] << 8) | ((uint64_t)ent[2] << 16) |
                            ((uint64_t)ent[3] << 24) | ((uint64_t)ent[4] << 32) |
                            ((uint64_t)ent[5] << 40) | ((uint64_t)ent[6] << 48) | ((uint64_t)ent[7] << 56);
                else
                    entry = ent[0] | (ent[1] << 8) | (ent[2] << 16) | (ent[3] << 24);

                if (entry == 0 || entry > 0x100000000ULL) break;
                auto ranges = m_bin->getExecutableRanges();
                bool ok = false;
                for (const auto& r : ranges)
                    if (entry >= r.first && entry < r.second) { ok = true; break; }
                if (ok) return table;
                else break;
            }
        }
    }
    return 0;
}

addr_t CFGBuilder::traceReg(const string& reg, const vector<Instruction>& block, size_t idx)
{
    if (idx >= block.size() || idx == 0) return 0;
    string cur = reg;

    for (int i = (int)idx - 1; i >= 0; i--)
    {
        const auto& inst = block[i];
        if (inst.mnemonic == "nop" || inst.mnemonic == "cmp" || inst.mnemonic == "test")
            continue;

        size_t comma = inst.operands.find(',');
        string dst = comma != string::npos ? inst.operands.substr(0, comma) : inst.operands;
        while (!dst.empty() && dst[0] == ' ') dst = dst.substr(1);
        while (!dst.empty() && dst.back() == ' ') dst.pop_back();

        if (inst.mnemonic == "mov" || inst.mnemonic == "lea")
        {
            if (dst == cur || dst == "r" + cur.substr(1) || regMatch(dst, cur))
            {
                if (inst.mnemonic == "mov" && !inst.operands.empty())
                {
                    size_t pos = inst.operands.rfind("0x");
                    if (pos != string::npos)
                    {
                        try { return stoull(inst.operands.substr(pos), nullptr, 16); }
                        catch (...) {}
                    }
                }
                if (inst.mnemonic == "lea")
                {
                    int64_t disp = m_dis->getRIPDisp(inst);
                    if (disp != 0) return inst.address + inst.size + disp;
                }
                if (comma != string::npos)
                {
                    string src = inst.operands.substr(comma + 1);
                    while (!src.empty() && src[0] == ' ') src = src.substr(1);
                    while (!src.empty() && src.back() == ' ') src.pop_back();
                    if (!src.empty() && src[0] != '0' && !isMemOp(src))
                    {
                        cur = src;
                        continue;
                    }
                }
                return 0;
            }
        }
        if (dst == cur || regMatch(dst, cur)) return 0;
    }
    return 0;
}

void CFGBuilder::recordIndirect(const Instruction& inst, addr_t target, bool ok, const string& name, const string& lib)
{
    IndirectTarget it;
    it.instruction_addr = inst.address;
    it.resolved_target = ok ? target : 0;
    it.resolved = ok;
    it.name = name;
    it.library = lib;
    m_targets.push_back(it);
}

BasicBlock CFGBuilder::disassembleBlock(addr_t start, addr_t& next)
{
    BasicBlock block;
    block.address = start;
    block.is_prolog = false;
    block.is_epilog = false;

    addr_t cur = start;
    auto ranges = m_bin->getExecutableRanges();

    auto inExec = [&](addr_t a) -> bool {
        for (const auto& r : ranges)
            if (a >= r.first && a < r.second) return true;
        return false;
    };

    for (size_t i = 0; i < 5000; ++i)
    {
        if (!inExec(cur)) { next = 0; break; }

        auto bytes = m_bin->readBytes(cur, 15);
        if (bytes.empty()) { next = 0; break; }

        auto insts = m_dis->disassemble(bytes, cur, 1);
        if (insts.empty()) { next = 0; break; }

        auto& inst = insts[0];
        block.instructions.push_back(inst);

        if (i <= 3 && !block.is_prolog && m_dis->isProlog(block.instructions))
            block.is_prolog = true;

        if (m_dis->isEpilog(inst))
            block.is_epilog = true;

        if (m_dis->isReturn(inst)) { next = 0; break; }
        if (m_dis->isTrap(inst)) { next = 0; break; }

        if (m_dis->isUnconditionalBranch(inst))
        {
            addr_t t = m_dis->getBranchTarget(inst);
            if (t && !m_dis->isIndirectBranch(inst))
                block.successors.push_back(t);
            else if (m_dis->isIndirectBranch(inst))
            {
                addr_t resolved = resolveGOT(inst);
                if (!resolved)
                {
                    resolved = resolveJumpTable(block.instructions, inst);
                    if (resolved)
                    {
                        for (int j = (int)block.instructions.size() - 2; j >= 0; j--)
                        {
                            if (block.instructions[j].mnemonic == "lea")
                            {
                                int64_t ld = m_dis->getRIPDisp(block.instructions[j]);
                                if (ld != 0)
                                {
                                    addr_t tbl = block.instructions[j].address + block.instructions[j].size + ld;
                                    for (int64_t e = 0; e < 64; e++)
                                    {
                                        auto eb = m_bin->readBytes(tbl + e * 8, 8);
                                        if (eb.size() < 4) break;
                                        uint64_t entry = 0;
                                        if (eb.size() >= 8)
                                            entry = eb[0] | ((uint64_t)eb[1] << 8) | ((uint64_t)eb[2] << 16) |
                                                    ((uint64_t)eb[3] << 24) | ((uint64_t)eb[4] << 32) |
                                                    ((uint64_t)eb[5] << 40) | ((uint64_t)eb[6] << 48) | ((uint64_t)eb[7] << 56);
                                        else
                                            entry = eb[0] | (eb[1] << 8) | (eb[2] << 16) | (eb[3] << 24);
                                        if (entry == 0 || entry > 0x100000000ULL) break;
                                        bool ok = false;
                                        for (const auto& r : ranges)
                                            if (entry >= r.first && entry < r.second) { ok = true; break; }
                                        if (ok) block.successors.push_back(entry);
                                        else break;
                                    }
                                    resolved = tbl;
                                    break;
                                }
                            }
                        }
                    }
                }
                if (!resolved)
                {
                    string reg = inst.operands;
                    while (!reg.empty() && reg[0] == ' ') reg = reg.substr(1);
                    resolved = traceReg(reg, block.instructions, block.instructions.size() - 1);
                    if (resolved)
                    {
                        block.successors.push_back(resolved);
                        recordIndirect(inst, resolved, true, "traced", "");
                    }
                    else
                    {
                        recordIndirect(inst, 0, false, "unresolved", "");
                    }
                }
                else
                {
                    auto git = m_got.find(resolved);
                    string nm = git != m_got.end() ? git->second.first : "";
                    string lb = git != m_got.end() ? git->second.second : "";
                    block.successors.push_back(resolved);
                    recordIndirect(inst, resolved, true, nm, lb);
                }
            }
            next = 0; break;
        }

        if (m_dis->isConditionalBranch(inst))
        {
            addr_t t = m_dis->getBranchTarget(inst);
            if (t && !m_dis->isIndirectBranch(inst))
                block.successors.push_back(t);
            block.successors.push_back(cur + inst.size);
            next = cur + inst.size; break;
        }

        if (m_dis->isCall(inst))
        {
            addr_t t = m_dis->getBranchTarget(inst);
            if (t && !m_dis->isIndirectBranch(inst))
            {
                m_funcs.insert(t);
                recordIndirect(inst, t, true, "", "");
            }
            else if (m_dis->isIndirectBranch(inst))
            {
                addr_t resolved = resolveGOT(inst);
                if (resolved)
                {
                    auto git = m_got.find(resolved);
                    string nm = git != m_got.end() ? git->second.first : "";
                    string lb = git != m_got.end() ? git->second.second : "";
                    auto gb = m_bin->readBytes(resolved, 8);
                    if (gb.size() >= 8)
                    {
                        uint64_t target = gb[0] | ((uint64_t)gb[1] << 8) |
                                          ((uint64_t)gb[2] << 16) | ((uint64_t)gb[3] << 24) |
                                          ((uint64_t)gb[4] << 32) | ((uint64_t)gb[5] << 40) |
                                          ((uint64_t)gb[6] << 48) | ((uint64_t)gb[7] << 56);
                        if (target && target < 0x100000000ULL) resolved = target;
                    }
                    m_funcs.insert(resolved);
                    recordIndirect(inst, resolved, true, nm, lb);
                }
                else
                {
                    string reg = inst.operands;
                    while (!reg.empty() && reg[0] == ' ') reg = reg.substr(1);
                    resolved = traceReg(reg, block.instructions, block.instructions.size() - 1);
                    if (resolved)
                    {
                        m_funcs.insert(resolved);
                        recordIndirect(inst, resolved, true, "traced", "");
                    }
                    else
                    {
                        recordIndirect(inst, 0, false, "unresolved", "");
                    }
                }
            }
        }

        cur += inst.size;
        next = cur;
    }

    return block;
}

set<addr_t> CFGBuilder::findPrologs()
{
    set<addr_t> cand;
    auto ranges = m_bin->getExecutableRanges();

    for (const auto& r : ranges)
    {
        addr_t a = r.first;
        while (a < r.second)
        {
            auto bytes = m_bin->readBytes(a, 15);
            if (bytes.empty()) { a += 1; continue; }
            auto insts = m_dis->disassemble(bytes, a, 1);
            if (insts.empty()) { a += 1; continue; }
            auto& inst = insts[0];

            if (inst.mnemonic == "endbr64")
            {
                auto nb = m_bin->readBytes(a + inst.size, 15);
                auto ni = m_dis->disassemble(nb, a + inst.size, 2);
                if (!ni.empty() && m_dis->isProlog(ni)) cand.insert(a);
            }
            else
            {
                vector<Instruction> tmp{inst};
                if (m_dis->isProlog(tmp)) cand.insert(a);
            }

            a += inst.size;
        }
    }

    return cand;
}

Function CFGBuilder::buildFunction(addr_t start, const string& name)
{
    Function func;
    func.address = start;
    func.name = name;

    queue<addr_t> q;
    q.push(start);

    while (!q.empty())
    {
        addr_t a = q.front(); q.pop();
        if (isVisited(a)) continue;
        markVisited(a);

        addr_t nxt = a;
        auto block = disassembleBlock(a, nxt);
        for (auto s : block.successors)
            if (!isVisited(s)) q.push(s);

        func.blocks.push_back(block);
    }

    return func;
}

CFG CFGBuilder::build()
{
    CFG cfg;
    cfg.binary_path = m_bin->getPath();
    cfg.entry_point = m_bin->getEntryPoint();

    switch (m_bin->getArch())
    {
        case Arch::X86: cfg.arch_str = "x86"; break;
        case Arch::X64: cfg.arch_str = "x86-64"; break;
        default: cfg.arch_str = "unknown";
    }
    switch (m_bin->getFormat())
    {
        case BinaryFormat::PE: cfg.format_str = "PE"; break;
        case BinaryFormat::ELF: cfg.format_str = "ELF"; break;
        default: cfg.format_str = "unknown";
    }

    cfg.imports = m_bin->getImportedFunctions();
    buildGotMap();

    m_prologs = findPrologs();

    set<addr_t> built;
    queue<addr_t> pending;

    auto process = [&](addr_t addr, const string& name)
    {
        if (built.count(addr)) return;
        built.insert(addr);

        auto func = buildFunction(addr, name);
        cfg.functions.push_back(func);

        for (const auto& b : func.blocks)
            for (const auto& inst : b.instructions)
                if (m_dis->isCall(inst))
                {
                    addr_t t = m_dis->getBranchTarget(inst);
                    if (t && !m_dis->isIndirectBranch(inst))
                        m_funcs.insert(t);
                }

        for (auto d : m_funcs)
            if (!built.count(d)) pending.push(d);
    };

    {
        addr_t ep = m_bin->getEntryPoint();
        string ep_name = "entry";
        for (const auto& exp : m_bin->getExportedFunctions())
            if (exp.first == ep) { ep_name = exp.second; break; }
        process(ep, ep_name);
    }

    for (const auto& exp : m_bin->getExportedFunctions())
        pending.push(exp.first);

    map<addr_t, string> imp_map;
    for (const auto& imp : m_bin->getImportedFunctions())
        imp_map[imp.address] = imp.name;

    while (!pending.empty())
    {
        addr_t a = pending.front(); pending.pop();
        if (built.count(a)) continue;

        string name;
        for (const auto& exp : m_bin->getExportedFunctions())
            if (exp.first == a) { name = exp.second; break; }
        if (name.empty())
        {
            auto it = imp_map.find(a);
            if (it != imp_map.end()) name = it->second;
        }

        built.insert(a);
        auto func = buildFunction(a, name);
        cfg.functions.push_back(func);

        for (const auto& b : func.blocks)
            for (const auto& inst : b.instructions)
                if (m_dis->isCall(inst))
                {
                    addr_t t = m_dis->getBranchTarget(inst);
                    if (t && !m_dis->isIndirectBranch(inst))
                        if (!built.count(t)) pending.push(t);
                }

        for (auto p : m_prologs)
            if (!built.count(p)) pending.push(p);
    }

    cfg.indirect_targets = m_targets;
    return cfg;
}