#include "types.hpp"
#include "loader/binary.hpp"
#include "cfg/builder.hpp"
#include "disasm/engine.hpp"

int main(int argc, char* argv[])
{
    if (argc < 2)
    {
        cerr << "usage: cfgrip [--subs-only] [--clean] <binary>" << endl;
        return 1;
    }

    bool subsOnly = false;
    bool clean = false;
    string path;

    for (int i = 1; i < argc; i++)
    {
        string arg = argv[i];
        if (arg == "--subs-only")
            subsOnly = true;
        else if (arg == "--clean")
            clean = true;
        else
            path = arg;
    }

    if (path.empty())
    {
        cerr << "usage: cfgrip [--subs-only] [--clean] <binary>" << endl;
        return 1;
    }

    auto binary = Binary::create(path);
    if (!binary)
    {
        cerr << "failed to load binary: " << path << endl;
        return 1;
    }

    cout << "format: ";
    switch (binary->getFormat())
    {
        case BinaryFormat::PE: cout << "PE"; break;
        case BinaryFormat::ELF: cout << "ELF"; break;
        default: cout << "unknown";
    }
    cout << endl;

    cout << "arch: ";
    switch (binary->getArch())
    {
        case Arch::X86: cout << "x86"; break;
        case Arch::X64: cout << "x86-64"; break;
        default: cout << "unknown";
    }
    cout << endl;

    cout << "entry: 0x" << hex << binary->getEntryPoint() << dec << endl;
    cout << "imports: " << binary->getImportedFunctions().size() << endl;

    for (const auto& imp : binary->getImportedFunctions())
    {
        cout << "  0x" << hex << imp.address << " " << imp.name;
        if (!imp.library.empty()) cout << " (" << imp.library << ")";
        cout << dec << endl;
    }

    Disassembler dis;
    if (!dis.init(binary->getArch()))
    {
        cerr << "failed to init capstone" << endl;
        return 1;
    }

    CFGBuilder builder(binary.get(), &dis, subsOnly);
    CFG cfg = builder.build();
    cfg.mode_str = subsOnly ? "subs-only" : "full";

    if (clean)
    {
        CFGBuilder::clean(cfg);
        cfg.mode_str += "+clean";
        cout << "clean: xrefs=" << cfg.xrefs.size() << " functions=" << cfg.functions.size() << endl;
    }

    cout << "functions: " << cfg.functions.size() << endl;
    cout << "indirect targets: " << cfg.indirect_targets.size() << endl;

    string suffix;
    if (subsOnly) suffix += ".subs";
    if (clean) suffix += ".clean";
    string out = path + suffix + ".cfg";
    ofstream f(out);
    if (f)
    {
        f << cfg.toJSON();
        f.close();
        cout << "cfg written to: " << out << endl;
    }

    return 0;
}