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

#include <fxos/view/assembly.h>
#include <fxos/view/util.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>

namespace FxOS {

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

static inline bool output(OperandOutput &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 renderOperand(AsmOperand const &op, u32 pc, int opsize,
    OperandOutput &out, ViewAssemblyOptions const &opts)
{
    out.push_back({{}, op.str()});

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

    if(op.kind() == AsmOperand::PcJump)
        type = PCJump;
    else if(op.kind() == AsmOperand::PcRel)
        type = PCRelative;
    else if(op.kind() == AsmOperand::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;
        if(!op.usesPCRelativeAddressing())
            return;

        u32 location = op.getPCRelativeTarget(pc);
        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;
    OperandOutput opout;

    /* 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");
    }

    /* Only show the raw data if instruction cannot be decoded */
    if(opts.showInstructionDetails)
        printf(" %08x:  %04x", pc, (i.inst ? i.inst->encoding() : i.opcode));

    if(!i.inst) {
        if(!opts.showInstructionDetails)
            printf(" %04x", i.inst->encoding());
        printf("\n");
        m_lastAddress = pc;
        return;
    }

    /* Mnemonic */
    std::string mnemonic = i.inst->mnemonic();
    std::string str = "   ";
    str += mnemonic;
    int spacing = i.inst->operandCount() ? 8 - mnemonic.size() : 0;
    printf("   %s%*s", mnemonic.c_str(), spacing, "");

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

        renderOperand(i.inst->operand(n), pc, i.inst->opsize(), opout, opts);

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

        opout.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();
    OperandOutput opout;
    u32 pc = ins.address();

    if(opts->showInstructionDetails)
        printf("    %08x:  %04x   ", pc, opcode.encoding());
    else
        printf("    ");

    /* Mnemonic */
    std::string mnemonic = opcode.mnemonic();
    std::string str = "";
    str += mnemonic;
    str += std::string(opcode.operandCount() ? 8 - mnemonic.size() : 0, ' ');

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

        renderOperand(opcode.operand(n), pc, opcode.opsize(), opout, *opts);

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

        opout.clear();
    }

    std::vector<std::pair<std::string, fmt::text_style>> comments;

    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(diff.target()
                != static_cast<int>(ProgramStateDiff::Target::None))
                comments.emplace_back(
                    diff.str(), fmt::fg(fmt::terminal_color::cyan));
        }
    }

    std::cout << str;

    if(!comments.empty()) {
        if(str.size() < 28)
            std::cout << std::string(28 - str.size(), ' ');

        bool first = true;
        for(auto &[c, style]: comments) {
            if(first)
                fmt::print(fmt::fg(fmt::color::gray), "# ");
            else
                fmt::print(", ");

            fmt::print(style, "{}", c);
            first = false;
        }
    }

    std::cout << '\n';
}

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

static void viewProgramState(ProgramState const &PS, std::string lineStart)
{
    ViewStringsOptions opts = {
        .maxColumns = 70,
        .lineStart = lineStart,
        .separator = ", ",
        .style = fmt::fg(fmt::terminal_color::cyan),
    };
    RelConstDomain RCD;
    std::cout << "\e[36m";
    viewStrings(std::views::iota(0, 16) | std::views::transform([&](int i) {
        return fmt::format("r{}:{}", i, PS.getRegister(i).str(false));
    }),
        opts);
    std::cout << "\e[0m";
}

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()];
            viewProgramState(block.entry, fmt_rgb("    | ", fmt::color::gray));
        }
    }

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