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

#include <fmt/core.h>

#include <fxos/disassembly.h>
#include <fxos/passes/cfg.h>
#include <fxos/passes/pcrel.h>
#include <fxos/passes/print.h>
#include <fxos/passes/syscall.h>
#include <fxos/util/Timer.h>
#include <fxos/util/log.h>

static void disassemble(Session &session, Disassembly &disasm,
    std::vector<std::string> const &passes, uint32_t address)
{
    for(auto pass: passes) {
        Timer timer;
        timer.start();

        bool ok;

        if(pass == "cfg") {
            CfgPass p(disasm);
            ok = p.analyzeAnonymousFunction(address);
        }
        else if(pass == "pcrel") {
            PcrelPass p(disasm);
            ok = p.analyzeAllInstructions();
        }
        else if(pass == "syscall") {
            OS *os = session.current_space->os_analysis();
            if(os) {
                SyscallPass p(disasm, os);
                ok = p.analyzeAllInstructions();
            }
        }
        else if(pass == "print") {
            PrintPass p(disasm);

            p.promote_pcjump_loc = PrintPass::Promote;
            p.promote_pcrel_loc = PrintPass::Promote;
            p.promote_pcrel_value = PrintPass::Promote;
            p.promote_syscall = PrintPass::Promote;
            p.promote_syscallname = PrintPass::Append;
            p.promote_symbol = PrintPass::Append;
            p.promote_pcaddr_loc = PrintPass::Promote;

            ok = p.analyzeAllInstructions();
        }
        else {
            FxOS_log(ERR, "unknown pass <%s>", pass);
            ok = false;
        }

        timer.stop();
        FxOS_log(LOG, "Finished pass <%s> in %s", pass, timer.format_time());

        if(!ok) {
            FxOS_log(ERR, "pass <%s> failed", pass);
            break;
        }
    }
}

//---
// d
//---

static uint32_t parse_d(Session &session, Parser &parser)
{
    if(!session.current_space)
        return 0;
    uint32_t address = session.current_space->cursor;

    if(!parser.at_end())
        address = parser.expr(session.current_space);

    parser.end();
    return address;
}

void _d(Session &session, uint32_t address)
{
    if(!session.current_space)
        return;
    FxOS::Disassembly disasm(*session.current_space);

    if(address & 1) {
        fmt::print("address 0x{:08x} is odd, starting at 0x{:08x}\n", address,
            address + 1);
        address++;
    }

    disassemble(session, disasm,
        {"cfg", "pcrel", /*"constprop",*/ "syscall", "print"}, address);
}

//---
// dr
//---

static Range parse_dr(Session &session, Parser &parser)
{
    Range range = parser.range(session.current_space);
    parser.end();
    return range;
}

void _dr(Session &session, Range range)
{
    if(!session.current_space)
        return;
    FxOS::Disassembly disasm(*session.current_space);

    if(range.start & 1) {
        fmt::print("address 0x{:08x} is odd, starting at 0x{:08x}\n",
            range.start, range.start + 1);
        range.start++;
    }
    if(range.end & 1) {
        fmt::print("address 0x{:08x} is odd, ending at 0x{:08x}\n", range.end,
            range.end - 1);
        range.end--;
    }

    if(range.start >= range.end)
        return;

    /* Load the block into memory */
    for(uint32_t pc = range.start; pc < range.end; pc += 2)
        disasm.getInstructionAt(pc, true);

    disassemble(
        session, disasm, {"pcrel", /*"constprop",*/ "syscall", "print"}, -1);
}

//---
// Command registration
//---

static ShellCommand _d_cmd(
    "d", [](Session &s, Parser &p) { _d(s, parse_d(s, p)); },
    [](Session &s, Parser &p) { parse_d(s, p); }, "Disassemble", R"(
d [<address>]

Disassembles code starting at the specified address, exploring branches until
function terminators, invalid instructions, or dynamically-computed jumps. The
default address is $ (the cursor of the current virtual space).

This command does not extend the virtual space's main disassembly. It reads
analysis results from the virtual space, but doesn't add new information. Try
as? to disassemble in the space's main disassembly.

The following disassembler passes are run:
 cfg         Explores the code reachable from the start address
 pcrel       Computes PC-relative addresses (eg mov.l, mova, bf, bra...)
 syscall     Annotates uses of syscall table entries with the syscall number
)");

static ShellCommand _dr_cmd(
    "dr", [](Session &s, Parser &p) { _dr(s, parse_dr(s, p)); },
    [](Session &s, Parser &p) { parse_dr(s, p); }, "Disassemble Range", R"(
dr [<range>]

Disassembles an explicit region of memory. This is similar to d, except that
the disassembled code is pre-loaded from the region instead of being explored
by the cfg pass. See d? for more information.

Like d, this command does not extend the virtual space's main disassembly.
)");