#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
//---

struct _d_args
{
    std::variant<long, Range> location;
};

static _d_args parse_d(Session &session, Parser &parser)
{
    _d_args args;

    if(!session.current_space)
        return {};

    args.location = parser.at_end()
                        ? session.current_space->cursor
                        : parser.expr_or_range(session.current_space);

    parser.end();
    return args;
}

void _d(Session &session, std::variant<long, Range> location)
{
    if(!session.current_space)
        return;
    FxOS::Disassembly disasm(*session.current_space);


    if(std::holds_alternative<Range>(location)) {
        Range range = std::get<Range>(location);
        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);
    }
    else {
        uint32_t address = std::get<long>(location);

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

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

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

static ShellCommand _d_cmd(
    "d",
    [](Session &s, Parser &p) {
        auto args = parse_d(s, p);
        _d(s, args.location);
    },
    [](Session &s, Parser &p) { parse_d(s, p); }, "Disassemble", R"(
d [<address|range>]

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
)");