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

#include <algorithm>
#include <fmt/core.h>
#include <fmt/chrono.h>
#include <fxos/util/log.h>

//---
// ib
//---

static void show_vspace(VirtualSpace const &s)
{
    fmt::print("  Region  Start         End         File\n  ");
    for(int i = 0; i < 70; i++)
        fmt::print("─");
    fmt::print("\n");

    if(s.bindings.size() == 0) {
        fmt::print("  (no bindings)\n");
        return;
    }
    for(auto &b: s.bindings) {
        MemoryRegion const *ref = MemoryRegion::region_for(b.region);
        fmt::print("  {:<7s} 0x{:08x} .. 0x{:08x}", (ref ? ref->name : ""),
            b.region.start, b.region.end);
        if(b.buffer.path != "")
            fmt::print("  {}", b.buffer.path);
        fmt::print("\n");
    }
}

static void show_binary_short(
    std::string const &name, bool current, Binary const &b)
{
    auto const &objects = b.objects();
    u32 total_size = 0;
    for(auto const &[_, obj]: objects)
        total_size += obj->size();

    if(current)
        fmt::print("* ");
    fmt::print(theme(11), "{}\n", name);
    fmt::print(
        "  {} objects (totaling {} bytes)\n", objects.size(), total_size);
    fmt::print("\n");
    show_vspace(b.vspace());
}


void _ib(Session &session)
{
    // TODO: Factor these errors into Session.
    Binary *b = session.currentBinary();
    if(!b) {
        FxOS_log(ERR, "No current binary!");
        return;
    }

    show_binary_short(session.currentBinaryName(), true, *b);
    // TODO: Show more binary information
}

//---
// ibs
//---

std::vector<std::string> parse_ibs(Session &, Parser &parser)
{
    std::vector<std::string> args;

    while(!parser.at_end())
        args.push_back(parser.symbol("binary_name"));

    parser.end();
    return args;
}

void _ibs(Session &session, std::vector<std::string> const &args)
{
    bool first = true;

    for(auto const &name: args) {
        if(!first)
            fmt::print("\n");
        else
            first = false;

        Binary *b = session.project().getBinary(name);
        if(b)
            show_binary_short(name, name == session.currentBinaryName(), *b);
        else
            FxOS_log(
                ERR, "No binary named “%s” in current project!", name.c_str());
    }

    if(!args.size()) {
        for(auto const &[name, b]: session.project().binaries()) {
            if(!first)
                fmt::print("\n");
            else
                first = false;

            show_binary_short(name, name == session.currentBinaryName(), b);
        }
    }
}

//---
// if
//---

struct _if_args
{
    std::vector<uint32_t> addresses;
};

static struct _if_args parse_if(Session &session, Parser &parser)
{
    _if_args args;

    while(!parser.at_end())
        args.addresses.push_back(parser.expr(session.currentBinary()));

    parser.end();
    return args;
}

void _if(Session &session, struct _if_args const &args)
{
    if(!session.currentBinary())
        return;
    Disassembly &disasm = session.currentBinary()->vspace().disasm;

    if(!args.addresses.size()) {
        fmt::print("{} functions\n", disasm.functions.size());
    }

    for(uint32_t address: args.addresses) {
        OldFunction *func = disasm.getFunctionAt(address);
        if(!func) {
            FxOS_log(ERR, "no function at 0x{:08x}", address);
            continue;
        }

        // TODO: Promote address to syscall, name, etc.
        fmt::print("0x{:08x}:\n", address);
        fmt::print("  callTargets:\n");

        auto &ct = func->callTargets;
        // TODO: Promote address to syscall, name, etc.
        for(uint32_t pc: ct)
            fmt::print("   0x{:08x}\n", pc);
        if(!ct.size())
            fmt::print("    (none)\n");
    }
}

//---
// io
//---

struct _io_args
{
    Binary *binary;
    std::vector<u32> addresses;
};

static struct _io_args parse_io(Session &session, Parser &parser)
{
    _io_args args;
    std::string binname = session.currentBinaryName();

    parser.option(
        "binary", [&binname](std::string const &value) { binname = value; });
    parser.accept_options();

    args.binary = session.project().getBinary(binname);
    if(!args.binary) {
        std::string msg = fmt::format("No binary “{}” in project!", binname);
        if(parser.completing())
            throw Parser::CompletionRequest("_error", msg);
        else
            FxOS_log(ERR, "%s", msg.c_str());
    }

    while(!parser.at_end())
        args.addresses.push_back(parser.expr(args.binary));

    parser.end();
    return args;
}

static void print_object(BinaryObject &obj)
{
    fmt::print("  0x{:08x} ({}) ", obj.address(), obj.size());
    if(obj.name() != "")
        fmt::print("{} ", obj.name());

    if(obj.isFunction())
        fmt::print("<function> (TODO)\n");
    else if(obj.isVariable())
        fmt::print("<variable> (TODO)\n");
    else if(obj.isMark())
        fmt::print("<mark> (TODO)\n");
    else
        fmt::print("<unknown object kind> o(x_x)o\n");

    if(obj.comment() != "")
        fmt::print("    {}\n", obj.comment());
}

void _io(Session &, Binary *binary, std::vector<u32> addresses)
{
    if(!addresses.size()) {
        /* This is always sorted */
        for(auto const &[_, object]: binary->objects())
            print_object(*object);
        return;
    }

    /* Find all objects related to provided addresses. Each object is indexed
       not by its own address but by the address of interest that caused it to
       be added to the list. */
    std::multimap<u32, BinaryObject *> objects;

    for(u32 address: addresses) {
        auto [it, end] = binary->objects().equal_range(address);
        for(; it != end; ++it)
            objects.insert({address, it->second.get()});

        auto const &covering = binary->objectsCovering(address);
        for(BinaryObject *obj: covering)
            objects.insert({address, obj});
    }

    for(auto [interest, obj]: objects) {
        if(obj->address() == interest)
            fmt::print("Object defined at {:08x}:\n", interest);
        else
            fmt::print("Object that covers {:08x}:\n", interest);

        print_object(*obj);
    }
}

//---
// ios
//---

static char const *info_str
    = "OS: type %s version %02d.%02d.%04d\n"
      "\n"
      "Header information:\n"
      "  Bootcode timestamp (DateA)                 (0x%000008x)  :  %s\n"
      "  Bootcode checksum                          (0x%000008x)  :  0x%08x\n"
      "  Serial number                              (0x%000008x)  :  %s\n"
      "  OS version                                 (0x%000008x)  :  %s\n";

static char const *footer_str
    = "\nFooter information:\n"
      "  Detected footer address                                  :  0x%08x\n"
      "  Langdata entries found                                   :  %d\n"
      "  OS date (DateO)                            (0x%000008x)  :  %s\n"
      "  OS checksum                                (0x%000008x)  :  0x%08x\n"
      "  Computed OS checksum                                     :  0x%08x\n";

static char const *syscall_str
    = "\nSyscall information:\n"
      "  Syscall table address                      (0x%000008x)  :  0x%08x\n"
      "  Entries that point to valid memory                       :  0x%x\n"
      "  First seemingly invalid entry                            :  0x%08x\n"
      "  Syscall entries outside ROM:\n";

static char const *syscall_nonrom_str = "    %%%04x -> %08x (%s memory)\n";

static std::string parse_ios(Session &, Parser &parser)
{
    std::string name = parser.at_end() ? "" : parser.symbol("binary_name");
    parser.end();
    return name;
}

void _ios(Session &session, std::string name)
{
    Binary *binary = session.currentBinary();
    if(name != "")
        binary = session.project().getBinary(name);
    if(!binary)
        return;

    VirtualSpace &vspace = binary->vspace();
    OS *os = binary->OSAnalysis();
    if(!os)
        throw CommandError("os analysis on '{}' failed", name);

    printf(info_str, (os->type == OS::FX ? "FX" : "CG"), os->version_major,
        os->version_minor, os->version_patch, os->bootcode_timestamp.address,
        os->bootcode_timestamp.value.c_str(), os->bootcode_checksum.address,
        os->bootcode_checksum, os->serial_number.address,
        os->serial_number.value.c_str(), os->version.address,
        os->version.value.c_str());

    if(os->footer == (uint32_t)-1) {
        printf("\nFooter could not be found.\n");
    }
    else {
        printf(footer_str, os->footer, os->langdata, os->timestamp.address,
            os->timestamp.value.c_str(), os->checksum.address, os->checksum,
            os->computed_checksum);
    }

    uint32_t syscall_table = os->syscall_table_address();
    uint32_t first_noncall
        = vspace.read_u32(syscall_table + 4 * os->syscall_count());

    printf(syscall_str, (os->type == OS::FX ? 0x8001007c : 0x8002007c),
        syscall_table, os->syscall_count(), first_noncall);

    int total = 0;
    for(int i = 0; i < os->syscall_count(); i++) {
        uint32_t e = os->syscall(i);
        MemoryRegion const *r = MemoryRegion::region_for(e);
        if(!r || r->name == "ROM" || r->name == "ROM_P2")
            continue;

        printf(syscall_nonrom_str, i, e, r->name.c_str());
        total++;
    }

    if(!total)
        printf("    (none)\n");
}

//---
// ip
//---

void _ip(Session &session)
{
    if(!session.hasProject())
        fmt::print("No current project o(x_x)o\n");
    else {
        Project &p = session.project();
        fmt::print("Current project: {}{}\n", p.name(), p.isDirty() ? "*" : "");
        if(p.path() != "")
            fmt::print("Path: {}\n", p.path());
        else
            fmt::print("No path yet (use ps to assign one)\n");

        fmt::print("Binaries:");
        if(!p.binaries().size())
            fmt::print(" (none)");
        for(auto const &[name, _]: p.binaries())
            fmt::print(" {}", name);
        fmt::print("\n");
    }

    auto const &entries = session.recentProjects().entries();
    fmt::print("\nRecent projects:\n");
    if(!entries.size())
        fmt::print("  (none)\n");
    for(auto &e: entries)
        fmt::print("  {} ({}, last used {:%c})\n", e.name, e.path,
            fmt::localtime(e.utime));

    if(session.legacySpaces.size() != 0) {
        fmt::print("\nLegacy vspaces:\n");

        for(auto const &[name, lvs]: session.legacySpaces) {
            fmt::print("  {} ({} symbols, {} syscalls)\n", name,
                lvs.symbols.size(), lvs.syscalls.size());
            for(auto const &b: lvs.bindings) {
                fmt::print("    ");
                for(auto const &r: b.regions)
                    fmt::print("0x{:08x} ", r.start);
                fmt::print("'{}'\n", b.path);
            }
        }
    }
}

//---
// isc
//---

struct _isc_args
{
    std::string binary_name;
    bool sort = false;
    std::vector<uint32_t> addresses;
};

static struct _isc_args parse_isc(Session &session, Parser &parser)
{
    _isc_args args;

    parser.option("sort",
        [&args](std::string const &value) { args.sort = (value == "true"); });
    parser.option("binary",
        [&args](std::string const &value) { args.binary_name = value; });

    parser.accept_options();

    Binary *binary = session.currentBinary();
    if(!args.binary_name.empty()) {
        binary = session.project().getBinary(args.binary_name);
        if(!binary) {
            std::string msg
                = format("No binary “%s” in project!", args.binary_name);
            if(parser.completing())
                throw Parser::CompletionRequest("_error", msg);
            else
                FxOS_log(ERR, "%s", msg.c_str());
        }
    }

    while(!parser.at_end())
        args.addresses.push_back(parser.expr(binary));

    parser.end();
    return args;
}

struct SyscallInfo
{
    uint32_t address;
    int id;
};

bool operator<(const SyscallInfo &left, const SyscallInfo &right)
{
    return (left.address < right.address) || (left.id < right.id);
}

void _isc(Session &session, std::string binary_name, bool sort,
    std::vector<uint32_t> addresses)
{
    Binary *binary = session.currentBinary();
    if(!binary) {
        FxOS_log(ERR, "No current binary!");
        return;
    }

    if(!binary_name.empty())
        binary = session.project().getBinary(binary_name);

    if(!binary) {
        FxOS_log(ERR, "No binary “%s” in project!", binary_name.c_str());
        return;
    }

    OS *os = binary->OSAnalysis();
    if(!os) {
        if(!binary_name.empty())
            FxOS_log(ERR, "OS analysis on '%s' failed", binary_name);
        else
            FxOS_log(ERR, "OS analysis failed");
        return;
    }

    if(!addresses.empty()) {
        if(sort)
            std::sort(&addresses[0], &addresses[addresses.size()]);

        for(uint32_t address: addresses) {
            int syscall = os->find_syscall(address);
            if(syscall == -1)
                continue;

            fmt::print(theme(3), "  0x{:08x}", address);
            fmt::print(theme(10), "  %{:04x}", syscall);
            fmt::print("\n");
        }

        return;
    }

    int total = os->syscall_count();
    auto info = std::make_unique<SyscallInfo[]>(total);

    for(int i = 0; i < total; i++)
        info[i] = (SyscallInfo) {.address = os->syscall(i), .id = i};

    if(sort)
        std::sort(&info[0], &info[total]);

    for(int i = 0; i < total; i++) {
        fmt::print(theme(3), "  0x{:08x}", info[i].address);
        fmt::print(theme(10), "  %{:04x}", info[i].id);
        fmt::print("\n");
    }
}

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

static ShellCommand _ib_cmd(
    "ib", [](Session &s, Parser &p) { p.end(), _ib(s); },
    [](Session &, Parser &p) { p.end(); }, "Info Binary (current)", R"(
ib

Prints detailed information about the currently-selected binary in the project.
)");

static ShellCommand _ibs_cmd(
    "ibs", [](Session &s, Parser &p) { _ibs(s, parse_ibs(s, p)); },
    [](Session &s, Parser &p) { parse_ibs(s, p); }, "Info Binaries", R"(
ibs [<binary>...]

Prints short information about named binaries in the current project. If no
argument is specified, prints information about all binaries.
)");

static ShellCommand _if_cmd(
    "if",
    [](Session &s, Parser &p) {
        auto args = parse_if(s, p);
        _if(s, args);
    },
    [](Session &s, Parser &p) { parse_if(s, p); }, "Info Function", R"(
if [<function>...]

Prints information about functions. Without arguments, prints vspace-level
statistics. With arguments, prints detailed function info.
)");

static ShellCommand _io_cmd(
    "io",
    [](Session &s, Parser &p) {
        auto args = parse_io(s, p);
        _io(s, args.binary, args.addresses);
    },
    [](Session &s, Parser &p) { parse_io(s, p); }, "Info Objects", R"(
io [binary=<binary_name>] [<expr>...]

Lists all objects in the specified binary (default: the current one). If
expressions are provided, instead list objects defined at these addresses or
that cover these addresses. With sort=true, sorts by increasing addresses.
)");

static ShellCommand _ios_cmd(
    "ios", [](Session &s, Parser &p) { _ios(s, parse_ios(s, p)); },
    [](Session &s, Parser &p) { parse_ios(s, p); }, "Info OS", R"(
ios [<binary_name>]

Prints information about the OS mapped in the named virtual space (defaults to
the current one). This usually requires an OS binary to be mapped to ROM.
)");

static ShellCommand _ip_cmd(
    "ip", [](Session &s, Parser &) { _ip(s); },
    [](Session &, Parser &p) { p.end(); }, "Info Projects", R"(
ip

Projects information about the currently-loaded project, and paths and names of
recent projects;
)");

static ShellCommand _isc_cmd(
    "isc",
    [](Session &s, Parser &p) {
        auto args = parse_isc(s, p);
        _isc(s, args.binary_name, args.sort, args.addresses);
    },
    [](Session &s, Parser &p) { parse_isc(s, p); }, "Info Syscalls", R"(
isc [sort=true] [binary=<binary_name>] [<address>...]

Prints the syscall table for the specified binary (default: the current one).
By default, syscalls are enumerated by syscall number. If sort=true is
specified, they are instead sorted by address.
)");