#include "shell.h"
#include "parser.h"
#include "commands.h"
#include "errors.h"
#include "util.h"

#include <fxos/disassembly.h>
#include <fxos/vspace.h>
#include <fxos/util/Timer.h>
#include <fxos/util/log.h>
#include <fxos/passes/cfg.h>
#include <fxos/passes/pcrel.h>
#include <fxos/passes/syscall.h>
#include <fmt/core.h>
#include <endian.h>

//---
// ad
//---

static void ad_disassemble_all(
    VirtualSpace &space, std::vector<uint32_t> const &addresses, bool force)
{
    int successes = 0, errors = 0;
    Timer timer;

    /* Analyze the CFGs of all functions */

    timer.start();
    CfgPass cfg_pass(space.disasm);

    /* We collect subfunction addresses while running the pass */
    for(int i = 0; i < (int)addresses.size(); i++) {
        uint32_t entry = addresses[i];
        printr("[cfg %d/%zu] Disassembling 0x%08x...", i + 1, addresses.size(),
            entry);
        if(!cfg_pass.exploreFunction(entry)) {
            FxOS_log(ERR, "while processing 0x%08x", entry);
            errors++;
            if(!force)
                return;
        }
        else {
            for(Claim const &c: cfg_pass.resultClaims())
                space.disasm.addExclusiveClaim(c);
            successes++;
        }
    }
    timer.stop();
    printf("\n");
    FxOS_log(LOG, "Finished pass <cfg> in %s", timer.format_time());

    /* Annotate all decoded instructions with pcrel/syscall
	   TODO: analyze only the functions, if possible */
    printr("[pcrel] Resolving PC-relative addressing modes...");
    timer.restart();
    PcrelPass pcrel_pass(space.disasm);
    if(!pcrel_pass.analyzeAllInstructions()) {
        errors++;
        if(!force)
            return;
    }
    timer.stop();
    printf("\n");
    FxOS_log(LOG, "Finished pass <pcrel> in %s", timer.format_time());

    printr("[syscall] Finding syscall references...");
    timer.restart();
    OS *os = space.os_analysis();
    if(os) {
        SyscallPass syscall_pass(space.disasm, os);
        if(!syscall_pass.analyzeAllInstructions()) {
            errors++;
            if(!force)
                return;
        }
    }
    timer.stop();
    printf("\n");
    FxOS_log(LOG, "Finished pass <syscall> in %s", timer.format_time());

    printf(
        "Successfully analyzed %d functions (%d errors)\n", successes, errors);

    /* TODO: Get subfunction addresses by abstract interpretation and keep
	   going recursively */
}

static std::vector<uint32_t> parse_ad(Session &session, Parser &parser)
{
    if(!session.current_space)
        return std::vector<uint32_t>();

    std::vector<uint32_t> addresses;
    do {
        addresses.push_back(parser.expr(session.current_space));
    }
    while(!parser.at_end());

    parser.end();
    return addresses;
}

void _ad(Session &session, std::vector<uint32_t> const &addresses)
{
    if(!session.current_space)
        return;

    VirtualSpace &space = *session.current_space;
    ad_disassemble_all(space, addresses, false);
}

//--
// ads
//---

static void parse_ads(Session &session, Parser &parser)
{
    if(!session.current_space)
        return;

    parser.end();
}

void _ads(Session &session)
{
    if(!session.current_space)
        return;

    VirtualSpace &space = *session.current_space;
    OS *os = space.os_analysis();

    if(!os) {
        printf("ads: OS analysis failed, cannot enumerate syscalls");
        return;
    }

    std::vector<uint32_t> addresses;
    for(int i = 0; i < os->syscall_count(); i++)
        addresses.push_back(os->syscall(i));

    ad_disassemble_all(space, addresses, true);
}

//---
// Command definitions
//---

static ShellCommand _ad_cmd(
    "ad",
    [](Session &s, Parser &p) {
        auto addresses = parse_ad(s, p);
        _ad(s, addresses);
    },
    [](Session &s, Parser &p) { parse_ad(s, p); }, "Analysis: Disassemble", R"(
ad [<addresses>...]

Disassemble the given set of addresses into the current virtual space's main
disassembly. The main disassembly is used for OS-wide features like cross-
reference search or call graphs.
)");

static ShellCommand _ads_cmd(
    "ads",
    [](Session &s, Parser &p) {
        parse_ads(s, p);
        _ads(s);
    },
    [](Session &s, Parser &p) { parse_ads(s, p); },
    "Analysis: Disassemble all Syscalls", R"(
ads

Disassembles all syscalls entries using ad, which stores the results in the
current virtual space's main disassembly. Unlike ad, this commands continues
even if some syscalls fail to disassemble.
)");