#pragma once
#include <cstdint>
#include <string>
#include <vector>
#include <sstream>
#include <iomanip>
#include <iostream>
#include <fstream>
#include <memory>


using namespace std;
using addr_t = uint64_t;

struct ImportEntry
{
    addr_t address;
    string name;
    string library;
};

struct IndirectTarget
{
    addr_t instruction_addr;
    addr_t resolved_target;
    string name;
    string library;
    bool resolved;
};

struct XrefEntry
{
    addr_t target;
    vector<pair<addr_t, string>> callers;
};

struct Instruction
{
    addr_t address;
    vector<uint8_t> bytes;
    string mnemonic;
    string operands;
    uint8_t size;
    int64_t stack_offset = 0;
};

struct BasicBlock
{
    addr_t address;
    vector<Instruction> instructions;
    vector<addr_t> successors;
    bool is_prolog;
    bool is_epilog;
};

struct Function
{
    addr_t address;
    addr_t end_address = 0;
    string name;
    vector<BasicBlock> blocks;
    bool is_thunk = false;
};

struct CFG
{
    string binary_path;
    string arch_str;
    string format_str;
    string mode_str;
    addr_t entry_point;
    vector<Function> functions;
    vector<ImportEntry> imports;
    vector<IndirectTarget> indirect_targets;
    vector<XrefEntry> xrefs;
    bool has_stack_offsets = false;

    static string escapeJSON(const string& s)
    {
        ostringstream os;
        for (char c : s)
        {
            switch (c)
            {
                case '"':  os << "\\\""; break;
                case '\\': os << "\\\\"; break;
                case '\n': os << "\\n"; break;
                case '\r': os << "\\r"; break;
                case '\t': os << "\\t"; break;
                default:   os << c;
            }
        }
        return os.str();
    }

    string toJSON() const
    {
        ostringstream os;
        os << "{\n";
        os << "  \"binary\": \"" << escapeJSON(binary_path) << "\",\n";
        os << "  \"mode\": \"" << escapeJSON(mode_str) << "\",\n";
        os << "  \"arch\": \"" << escapeJSON(arch_str) << "\",\n";
        os << "  \"format\": \"" << escapeJSON(format_str) << "\",\n";
        os << "  \"entry_point\": \"0x" << hex << entry_point << "\",\n";
        os << "  \"imports\": [\n";
        for (size_t i = 0; i < imports.size(); ++i) {
            os << "    {\n";
            os << "      \"address\": \"0x" << hex << imports[i].address << "\",\n";
            os << "      \"name\": \"" << escapeJSON(imports[i].name) << "\",\n";
            os << "      \"library\": \"" << escapeJSON(imports[i].library) << "\"\n";
            os << "    }";
            if (i < imports.size() - 1) os << ",";
            os << "\n";
        }
        os << "  ],\n";

        os << "  \"indirect_targets\": [\n";
        for (size_t i = 0; i < indirect_targets.size(); ++i) {
            const auto& it = indirect_targets[i];
            os << "    {\n";
            os << "      \"instruction\": \"0x" << hex << it.instruction_addr << "\",\n";
            if (it.resolved) {
                os << "      \"resolved_target\": \"0x" << it.resolved_target << "\",\n";
            } else {
                os << "      \"resolved_target\": null,\n";
            }
            os << "      \"name\": \"" << escapeJSON(it.name) << "\",\n";
            os << "      \"library\": \"" << escapeJSON(it.library) << "\"\n";
            os << "    }";
            if (i < indirect_targets.size() - 1) os << ",";
            os << "\n";
        }
        os << "  ],\n";

        os << "  \"functions\": [\n";
        for (size_t i = 0; i < functions.size(); ++i) {
            const auto& f = functions[i];
            os << "    {\n";
            os << "      \"address\": \"0x" << hex << f.address << "\",\n";
            if (f.end_address)
                os << "      \"end_address\": \"0x" << hex << f.end_address << "\",\n";
            os << "      \"name\": \"" << escapeJSON(f.name) << "\",\n";
            if (f.is_thunk)
                os << "      \"is_thunk\": true,\n";
            os << "      \"blocks\": [\n";
            for (size_t j = 0; j < f.blocks.size(); ++j) {
                const auto& b = f.blocks[j];
                os << "        {\n";
                os << "          \"address\": \"0x" << hex << b.address << "\",\n";
                os << "          \"size\": " << dec << b.instructions.size() << ",\n";
                os << "          \"is_prolog\": " << (b.is_prolog ? "true" : "false") << ",\n";
                os << "          \"is_epilog\": " << (b.is_epilog ? "true" : "false") << ",\n";
                os << "          \"instructions\": [\n";
                for (size_t k = 0; k < b.instructions.size(); ++k) {
                    const auto& inst = b.instructions[k];
                    os << "            {\n";
                    os << "              \"address\": \"0x" << hex << inst.address << "\",\n";
                    os << "              \"size\": " << dec << (int)inst.size << ",\n";
                    os << "              \"mnemonic\": \"" << escapeJSON(inst.mnemonic) << "\",\n";
                    os << "              \"operands\": \"" << escapeJSON(inst.operands) << "\"";
                    if (has_stack_offsets) {
                        os << ",\n";
                        os << "              \"stack_offset\": " << dec << inst.stack_offset;
                    }
                    os << "\n";
                    os << "            }";
                    if (k < b.instructions.size() - 1) os << ",";
                    os << "\n";
                }
                os << "          ],\n";
                os << "          \"successors\": [\n";
                for (size_t k = 0; k < b.successors.size(); ++k) {
                    os << "            \"0x" << hex << b.successors[k] << "\"";
                    if (k < b.successors.size() - 1) os << ",";
                    os << "\n";
                }
                os << "          ]\n";
                os << "        }";
                if (j < f.blocks.size() - 1) os << ",";
                os << "\n";
            }
            os << "      ]\n";
            os << "    }";
            if (i < functions.size() - 1) os << ",";
            os << "\n";
        }
        if (!xrefs.empty())
        {
            os << "  ],\n";
            os << "  \"xrefs\": [\n";
            for (size_t i = 0; i < xrefs.size(); ++i) {
                const auto& x = xrefs[i];
                os << "    {\n";
                os << "      \"target\": \"0x" << hex << x.target << "\",\n";
                os << "      \"callers\": [\n";
                for (size_t j = 0; j < x.callers.size(); ++j) {
                    os << "        {\n";
                    os << "          \"address\": \"0x" << hex << x.callers[j].first << "\",\n";
                    os << "          \"type\": \"" << escapeJSON(x.callers[j].second) << "\"\n";
                    os << "        }";
                    if (j < x.callers.size() - 1) os << ",";
                    os << "\n";
                }
                os << "      ]\n";
                os << "    }";
                if (i < xrefs.size() - 1) os << ",";
                os << "\n";
            }
            os << "  ]\n";
        }
        else
        {
            os << "  ]\n";
        }
        os << "}\n";
        return os.str();
    }
};