//---------------------------------------------------------------------------//
//  1100101 |_ mov #0, r4         __                                         //
//     11   |_ <0xb380 %5c4>     / _|_ _____ ___                             //
//     0110 |_ 3.50 -> 3.60     |  _\ \ / _ (_-<                             //
//          |_ base# + offset   |_| /_\_\___/__/                             //
//---------------------------------------------------------------------------//

#include <fxos/view/assembly.h>
#include <fxos/analysis.h>
#include <fxos/binary.h>
#include <fxos/function.h>
#include <fxos/util/format.h>
#include <fxos/util/Queue.h>
#include <vector>
#include <cstdio>
#include <cstring>
#include <numeric>
#include <fmt/color.h>

namespace FxOS {

/* Output for a single argument, which consists of one or more text segments
   each with their own text style. */
using ArgumentOutput = std::vector<std::pair<fmt::text_style, std::string>>;

static inline bool output(ArgumentOutput &out, ViewAssemblyOptions::Promotion p,
    fmt::text_style style, std::string str)
{
    if(p == ViewAssemblyOptions::Never)
        return true;
    if(p == ViewAssemblyOptions::Promote)
        out.pop_back();

    out.push_back({style, std::move(str)});
    return false;
}

// TODO: Take advantage of Instruction's info
static void renderArgument(AsmArgument const &arg, u32 pc, int opsize,
    ArgumentOutput &out, ViewAssemblyOptions const &opts)
{
    out.push_back({{}, arg.str()});

    // clang-format off
    enum { None, PCJump, PCRelative, PCAddr, Location, Constant, SyscallNumber,
           ObjectName }
        type = None;
    // clang-format on

    if(arg.kind == AsmArgument::PcJump)
        type = PCJump;
    else if(arg.kind == AsmArgument::PcRel)
        type = PCRelative;
    else if(arg.kind == AsmArgument::PcAddr)
        type = PCAddr;

    u32 value = 0;
    bool hasValue = false;
    OS *os = opts.binary ? opts.binary->OSAnalysis() : nullptr;
    VirtualSpace *vspace = opts.binary ? &opts.binary->vspace() : nullptr;

    if(type == PCJump || type == PCRelative || type == PCAddr) {
        auto p = (type == PCJump)   ? opts.promotions.PCJump_to_Location
                 : (type == PCAddr) ? opts.promotions.PCAddr_to_Location
                                    : opts.promotions.PCRelative_to_Location;
        u32 location = arg.getPCRelativeTarget(pc, opsize);
        if(location == 0xffffffff)
            return;
        if(output(out, p, {}, format("<%08x>", location)))
            return;
        type = (type == PCRelative) ? Location : Constant;

        // TODO: Check that this is a read operation!
        if(opsize != 0 && vspace && vspace->covers(location, opsize)) {
            if(opsize == 1) {
                value = vspace->read_i8(location);
                hasValue = true;
            }
            if(opsize == 2) {
                value = vspace->read_i16(location);
                hasValue = true;
            }
            if(opsize == 4) {
                value = vspace->read_i32(location);
                hasValue = true;
            }
        }
    }

    if(type == Location) {
        auto p = opts.promotions.ReadLocation_to_Constant;
        auto rc = RelConstDomain().constant(value);
        if(!hasValue || output(out, p, {}, rc.str()))
            return;
        type = Constant;
    }

    /* Promote to object name first if available... */
    if(type == Constant && hasValue && opts.binary) {
        auto p = opts.promotions.Constant_to_ObjectName;
        BinaryObject *obj = opts.binary->objectAt(value);

        if(obj) {
            if(output(out, p, {}, obj->name()))
                return;
            type = ObjectName;
        }
    }
    /* ... or, as a default, a syscall number */
    if(type == Constant && hasValue && os) {
        int syscall_id = os->find_syscall(value);
        if(syscall_id >= 0) {
            auto p = opts.promotions.Constant_to_SyscallNumber;
            if(output(out, p, {}, format("%%%04x", syscall_id)))
                return;
            type = SyscallNumber;
        }
    }
}

//=== Legacy-style instruction printer ===//

static void doOldInst(u32 pc, OldInstruction &i,
    ViewAssemblyOptions const &opts, u32 &m_lastAddress)
{
    OS *os = opts.binary ? opts.binary->OSAnalysis() : nullptr;
    ArgumentOutput argout;

    /* Ellipsis if there is a gap since last instruction */
    if(m_lastAddress + 1 != 0 && pc != m_lastAddress + 2)
        printf(" ...\n");

    /* Preliminary syscall number */
    int syscall_id;
    if(os && (syscall_id = os->find_syscall(pc)) >= 0) {
        printf("\n<%%%04x", syscall_id);
        BinaryObject *obj = opts.binary ? opts.binary->objectAt(pc) : nullptr;
        if(obj)
            printf(" %s", obj->name().c_str());
        printf(">\n");
    }

    /* Raw data if instruction cannot be decoded */
    printf(" %08x:  %04x", pc, (i.inst ? i.inst->encoding : i.opcode));
    if(!i.inst) {
        printf("\n");
        m_lastAddress = pc;
        return;
    }

    /* Mnemonic */
    static char const *suffixes[5] = {"", ".b", ".w", "", ".l"};
    char const *suffix = suffixes[(i.inst->opsize <= 4) ? i.inst->opsize : 0];

    int spacing = i.inst->arg_count
                      ? 8 - strlen(i.inst->mnemonic()) - strlen(suffix)
                      : 0;
    printf("   %s%s%*s", i.inst->mnemonic(), suffix, spacing, "");

    /* Arguments */
    for(size_t n = 0; n < i.inst->arg_count; n++) {
        if(n)
            printf(", ");

        renderArgument(i.inst->args[n], pc, i.inst->opsize, argout, opts);

        for(size_t i = 0; i < argout.size(); i++) {
            if(i != 0)
                printf(" ");
            printf("%s", argout[i].second.c_str());
        }

        argout.clear();
    }

    printf("\n");
    m_lastAddress = pc;
}

void viewAssemblyLegacyRegion(
    Binary &binary, MemoryRegion r, ViewAssemblyOptions *opts_ptr)
{
    ViewAssemblyOptions opts;
    if(opts_ptr)
        opts = *opts_ptr;
    opts.binary = &binary;

    u32 lastAddress = 0xffffffff;

    for(u32 pc = r.start & -2; pc <= r.end; pc += 2) {
        OldInstruction *i = binary.vspace().disasm.getInstructionAt(pc, true);
        if(i != nullptr)
            doOldInst(pc, *i, opts, lastAddress);
    }
}

void viewAssemblyLegacyAddress(
    Binary &binary, u32 pc, ViewAssemblyOptions *opts_ptr)
{
    ViewAssemblyOptions opts;
    if(opts_ptr)
        opts = *opts_ptr;
    opts.binary = &binary;

    u32 lastAddress = 0xffffffff;

    Queue<u32> queue;
    queue.enqueue(pc);

    while(!queue.empty()) {
        u32 pc = queue.pop();
        OldInstruction *i = binary.vspace().disasm.getInstructionAt(pc, true);
        if(i == nullptr)
            continue;

        /* Enqueue successors */
        if(!i->terminal && !i->jump)
            queue.enqueue(pc + 2);
        if(i->jump || i->condjump)
            queue.enqueue(i->jmptarget);
    }

    /* Print explored instructions in increasing order of addresses */
    for(u32 pc: queue.seen) {
        OldInstruction *i = binary.vspace().disasm.getInstructionAt(pc, false);
        if(i)
            doOldInst(pc, *i, opts, lastAddress);
    }
}

//=== Binary-API assembly printer ===//

static ViewAssemblyOptions defaultOptions {};

void viewAssemblyInstruction(Instruction const &ins, ViewAssemblyOptions *opts)
{
    opts = opts ? opts : &defaultOptions;

    AsmInstruction opcode = ins.opcode();
    ArgumentOutput argout;
    u32 pc = ins.address();

    printf("    %08x:  %04x", pc, opcode.encoding);

    /* Mnemonic */
    static char const *suffixes[5] = {"", ".b", ".w", "", ".l"};
    char const *suffix = suffixes[(opcode.opsize <= 4) ? opcode.opsize : 0];

    int spacing
        = opcode.arg_count ? 8 - strlen(opcode.mnemonic()) - strlen(suffix) : 0;
    std::string str = "   ";
    str += opcode.mnemonic();
    str += suffix;
    str += std::string(spacing, ' ');

    /* Arguments */
    for(size_t n = 0; n < opcode.arg_count; n++) {
        if(n)
            str += ", ";

        renderArgument(opcode.args[n], pc, opcode.opsize, argout, *opts);

        for(size_t i = 0; i < argout.size(); i++) {
            if(i != 0)
                str += " ";
            str += argout[i].second;
        }

        argout.clear();
    }

    if(opts->printFunctionAnalysis) {
        auto *an = ins.parentFunction().getAnalysis();

        if(an) {
            auto &block = an->blocks[ins.parentBlock().blockIndex()];
            ProgramStateDiff const &diff = block.diffs[ins.indexInBlock()];

            if(str.size() < 32)
                str += std::string(32 - str.size(), ' ');
            str += "# ";
            str += diff.str();
        }
    }

    fmt::print("{}\n", str);
}

static std::string objectsAt(
    Binary const &binary, u32 address, BinaryObject const *except = nullptr)
{
    std::vector<std::string> objects;

    OS *os = binary.OSAnalysis();
    if(os) {
        int index = os->find_syscall(address);
        if(index >= 0)
            objects.push_back(fmt::format("%{:04x}", index));
    }

    int unnamed = 0;
    for(BinaryObject const *obj: binary.objectsAt(address)) {
        if(obj == except)
            continue;
        if(obj->name() == "")
            unnamed++;
        else
            objects.push_back(obj->name());
    }

    if(unnamed > 0)
        objects.push_back(fmt::format("+{}", unnamed));

    return std::accumulate(objects.begin(), objects.end(), std::string {},
        [](auto l, auto const r) {
            return std::move(l) + (l.empty() ? "" : " ") + r;
        });
}

void viewAssemblyBasicBlock(BasicBlock const &bb, ViewAssemblyOptions *opts)
{
    opts = opts ? opts : &defaultOptions;

    printf("  bb.%08x", bb.address());
    if(bb.address() != bb.parentFunction().address()) {
        std::string others = objectsAt(bb.parentBinary(), bb.address());
        if(others != "")
            printf(" (%s)", others.c_str());
    }
    printf(":\n");

    if(opts->basicBlockDetails) {
        printf("    Successors:");
        for(u32 succ: bb.successorsByAddress())
            printf(" bb.%08x", succ);
        if(bb.successorCount() == 0)
            printf(" (none)");
        printf("\n");

        printf("    Predecessors:");
        for(u32 succ: bb.predecessorsByAddress())
            printf(" bb.%08x", succ);
        if(bb.predecessorCount() == 0)
            printf(" (none)");
        printf("\n");

        printf("    Flags:");
        if(bb.isEntryBlock())
            printf(" IsEntryBlock");
        if(bb.isTerminator())
            printf(" IsTerminator");
        if(bb.hasDelaySlot())
            printf(" HasDelaySlot");
        if(bb.hasNoTerminator())
            printf(" NoTerminator");
        if(!(bb.getFlags() & BasicBlock::ValidFlags))
            printf(" (none)");
        printf("\n");
    }

    if(opts->printFunctionAnalysis) {
        auto *an = bb.parentFunction().getAnalysis();
        if(an) {
            auto &block = an->blocks[bb.blockIndex()];
            printf("    Entry state:\n");
            fmt::print("{}", block.entry.str(6));
        }
    }

    for(Instruction const &ins: bb)
        viewAssemblyInstruction(ins, opts);

    printf("\n");
}

void viewAssemblyFunction(Function const &fun, ViewAssemblyOptions *opts)
{
    opts = opts ? opts : &defaultOptions;

    /* Note that the Function constructor sets "fun.%08x" as name by default */
    if(fun.name() != "")
        printf("%s", fun.name().c_str());
    else
        printf("fun.%08x", fun.address());

    std::string others = objectsAt(fun.parentBinary(), fun.address(), &fun);
    if(others != "")
        printf(" (%s)", others.c_str());
    printf(":\n");

    for(BasicBlock const &bb: fun)
        viewAssemblyBasicBlock(bb, opts);
}

} /* namespace FxOS */