#include "shim/trace/emitter.h"

#include <cmath>
#include <cstdarg>
#include <cstdio>
#include <cstdlib>
#include <cstring>

#include "shim/trace/sha256.h"

namespace shim::trace {

// ---- modes ----------------------------------------------------------------------------------

const char* mode_name(Mode m) {
    switch (m) {
        case Mode::Off: return "off";
        case Mode::Trace: return "trace";
        case Mode::Compare: return "compare";
        case Mode::Replace: return "replace";
    }
    return "off";
}

bool parse_mode(const char* text, Mode& out) {
    if (!text) return false;
    if (std::strcmp(text, "off") == 0) { out = Mode::Off; return true; }
    if (std::strcmp(text, "trace") == 0) { out = Mode::Trace; return true; }
    if (std::strcmp(text, "compare") == 0) { out = Mode::Compare; return true; }
    if (std::strcmp(text, "replace") == 0) { out = Mode::Replace; return true; }
    return false;
}

// ---- typed values -----------------------------------------------------------------------------

const char* Tv::type_name() const {
    switch (kind) {
        case Bool: return "bool";
        case Int: return itype;
        case Float: return f32 ? "f32" : "f64";
        case Str: return "str";
        case WStr: return "wstr";
        case Ptr: return "ptr";
        case Enum: return "enum";
        case Null: return "null";
        case Bytes: return "bytes";
        case List: return "list";
        case Set: return "set";
        case Struct: return "struct";
        case Json: return "json";
    }
    return "null";
}

Tv& Tv::add(const char* key, Tv v) {
    keys.emplace_back(key);
    items.push_back(std::move(v));
    return *this;
}

namespace tv {

namespace {
Tv make_int(const char* t, long long v) {
    Tv r;
    r.kind = Tv::Int;
    r.itype = t;
    r.i = v;
    return r;
}
Tv make_uint(const char* t, unsigned long long v) {
    Tv r;
    r.kind = Tv::Int;
    r.itype = t;
    r.is_unsigned = true;
    r.u = v;
    return r;
}
}  // namespace

Tv boolean(bool v) { Tv r; r.kind = Tv::Bool; r.b = v; return r; }
Tv i8(std::int8_t v) { return make_int("i8", v); }
Tv i16(std::int16_t v) { return make_int("i16", v); }
Tv i32(std::int32_t v) { return make_int("i32", v); }
Tv i64(std::int64_t v) { return make_int("i64", v); }
Tv u8(std::uint8_t v) { return make_uint("u8", v); }
Tv u16(std::uint16_t v) { return make_uint("u16", v); }
Tv u32(std::uint32_t v) { return make_uint("u32", v); }
Tv u64(std::uint64_t v) { return make_uint("u64", v); }
Tv f32(float v) { Tv r; r.kind = Tv::Float; r.f32 = true; r.f = v; return r; }
Tv f64(double v) { Tv r; r.kind = Tv::Float; r.f = v; return r; }

Tv str(const char* s) {
    if (!s) return ptr(static_cast<std::uintptr_t>(0));
    return str(s, std::strlen(s));
}
Tv str(const void* bytes, std::size_t n) {
    Tv r;
    r.kind = Tv::Str;
    r.s.assign(static_cast<const char*>(bytes), n);
    return r;
}
Tv wstr(const std::uint16_t* units, std::size_t n) {
    Tv r;
    r.kind = Tv::WStr;
    r.w.assign(units, units + n);
    return r;
}
Tv wstr(const wchar_t* s) {
    if (!s) return ptr(static_cast<std::uintptr_t>(0));
    Tv r;
    r.kind = Tv::WStr;
    for (; *s; ++s) r.w.push_back(static_cast<std::uint16_t>(*s));
    return r;
}
Tv ptr(const void* p) { return ptr(reinterpret_cast<std::uintptr_t>(p)); }
Tv ptr(std::uintptr_t p) { Tv r; r.kind = Tv::Ptr; r.p = p; return r; }
Tv enum_(long long v, const char* symbolic) {
    Tv r;
    r.kind = Tv::Enum;
    r.i = v;
    if (symbolic) r.s = symbolic;
    return r;
}
Tv null() { return Tv(); }

Tv bytes(const void* data, std::size_t n, std::size_t inline_max) {
    Tv r;
    r.kind = Tv::Bytes;
    r.blen = n;
    Sha256::digest_hex(data, n, r.sha);
    const std::uint8_t* p = static_cast<const std::uint8_t*>(data);
    if (n <= inline_max) {
        r.inline_full = true;
        r.data.assign(p, p + n);
    } else {
        r.data.assign(p, p + 32);
    }
    return r;
}
Tv list(std::vector<Tv> items) { Tv r; r.kind = Tv::List; r.items = std::move(items); return r; }
Tv set(std::vector<Tv> items) { Tv r; r.kind = Tv::Set; r.items = std::move(items); return r; }
Tv struct_() { Tv r; r.kind = Tv::Struct; return r; }
Tv json(const char* canonical_text) { Tv r; r.kind = Tv::Json; r.s = canonical_text; return r; }

}  // namespace tv

// ---- output buffer --------------------------------------------------------------------------------

Buf::Buf() { grow(4096); }
Buf::~Buf() { std::free(data_); }

bool Buf::grow(std::size_t need) {
    if (need <= cap_) return true;
    std::size_t ncap = cap_ ? cap_ : 4096;
    while (ncap < need) ncap *= 2;
    char* nd = static_cast<char*>(std::realloc(data_, ncap));
    if (!nd) {
        overflow_ = true;
        return false;
    }
    data_ = nd;
    cap_ = ncap;
    return true;
}

void Buf::put(char c) {
    if (len_ + 1 > cap_ && !grow(len_ + 1)) return;
    data_[len_++] = c;
}

void Buf::put(const void* bytes, std::size_t n) {
    if (!n) return;
    if (len_ + n > cap_ && !grow(len_ + n)) return;
    std::memcpy(data_ + len_, bytes, n);
    len_ += n;
}

void Buf::puts(const char* s) { put(s, std::strlen(s)); }

void Buf::printf(const char* fmt, ...) {
    char tmp[128];
    va_list ap;
    va_start(ap, fmt);
    const int n = std::vsnprintf(tmp, sizeof tmp, fmt, ap);
    va_end(ap);
    if (n < 0) {
        overflow_ = true;
        return;
    }
    if (static_cast<std::size_t>(n) < sizeof tmp) {
        put(tmp, static_cast<std::size_t>(n));
        return;
    }
    if (!grow(len_ + static_cast<std::size_t>(n) + 1)) return;
    va_start(ap, fmt);
    std::vsnprintf(data_ + len_, static_cast<std::size_t>(n) + 1, fmt, ap);
    va_end(ap);
    len_ += static_cast<std::size_t>(n);
}

void Buf::clear() {
    len_ = 0;
    overflow_ = false;
}

// ---- escaping / numbers -------------------------------------------------------------------------

namespace {

const char* const kHex = "0123456789abcdef";

inline void put_u_escape(Buf& out, unsigned unit) {
    char e[6] = {'\\', 'u', kHex[(unit >> 12) & 15], kHex[(unit >> 8) & 15], kHex[(unit >> 4) & 15], kHex[unit & 15]};
    out.put(e, 6);
}

// One unit of either string flavour (bytes for str, UTF-16 units for wstr): mkfixture.esc().
inline void put_unit(Buf& out, unsigned unit) {
    switch (unit) {
        case '"': out.put("\\\"", 2); return;
        case '\\': out.put("\\\\", 2); return;
        case '\n': out.put("\\n", 2); return;
        case '\r': out.put("\\r", 2); return;
        case '\t': out.put("\\t", 2); return;
        default:
            if (unit < 0x20 || unit >= 0x7f) put_u_escape(out, unit);
            else out.put(static_cast<char>(unit));
    }
}

void put_hex(Buf& out, const std::uint8_t* p, std::size_t n) {
    for (std::size_t i = 0; i < n; ++i) {
        out.put(kHex[p[i] >> 4]);
        out.put(kHex[p[i] & 15]);
    }
}

const long long kTwo53 = 1LL << 53;

}  // namespace

void emit_esc(Buf& out, const char* bytes, std::size_t n) {
    out.put('"');
    for (std::size_t i = 0; i < n; ++i) put_unit(out, static_cast<unsigned char>(bytes[i]));
    out.put('"');
}

void emit_esc(Buf& out, const std::string& s) { emit_esc(out, s.data(), s.size()); }

void emit_esc_w(Buf& out, const std::uint16_t* units, std::size_t n) {
    out.put('"');
    for (std::size_t i = 0; i < n; ++i) put_unit(out, units[i]);
    out.put('"');
}

// Integers above 2^53 in magnitude go out as decimal strings (section 3): mkfixture.u64().
void emit_int(Buf& out, long long v) {
    const bool quote = v >= kTwo53 || v <= -kTwo53;
    if (quote) out.put('"');
    out.printf("%lld", v);
    if (quote) out.put('"');
}

void emit_uint(Buf& out, unsigned long long v) {
    const bool quote = v >= static_cast<unsigned long long>(kTwo53);
    if (quote) out.put('"');
    out.printf("%llu", v);
    if (quote) out.put('"');
}

// mkfixture.fmt_num() for floats.
void emit_float(Buf& out, double v, bool f32) {
    if (std::isnan(v)) { out.puts("\"nan\""); return; }
    if (std::isinf(v)) { out.puts(v > 0 ? "\"inf\"" : "\"-inf\""); return; }
    out.printf(f32 ? "%.9g" : "%.17g", v);
}

// ---- typed values ----------------------------------------------------------------------------------

// mkfixture.emit_tv(): {"t":..., ("n","sha256","hex"|"head") | "v":... [,"name":...]} [,"n":...]
void emit_tv(Buf& out, const Tv& v) {
    out.puts("{\"t\":");
    emit_esc(out, v.type_name(), std::strlen(v.type_name()));
    if (v.kind == Tv::Bytes) {
        out.printf(",\"n\":%llu,\"sha256\":\"", static_cast<unsigned long long>(v.blen));
        out.puts(v.sha);
        out.puts(v.inline_full ? "\",\"hex\":\"" : "\",\"head\":\"");
        put_hex(out, v.data.data(), v.data.size());
        out.puts("\"}");
        return;
    }
    out.puts(",\"v\":");
    switch (v.kind) {
        case Tv::Bool: out.puts(v.b ? "true" : "false"); break;
        case Tv::Int:
            if (v.is_unsigned) emit_uint(out, v.u);
            else emit_int(out, v.i);
            break;
        case Tv::Float: emit_float(out, v.f, v.f32); break;
        case Tv::Str: emit_esc(out, v.s); break;
        case Tv::WStr: emit_esc_w(out, v.w.data(), v.w.size()); break;
        case Tv::Ptr: out.printf("\"0x%08llx\"", static_cast<unsigned long long>(v.p)); break;
        case Tv::Enum: out.printf("%lld", v.i); break;
        case Tv::Null: out.puts("null"); break;
        case Tv::List:
        case Tv::Set:
            out.put('[');
            for (std::size_t i = 0; i < v.items.size(); ++i) {
                if (i) out.put(',');
                emit_tv(out, v.items[i]);
            }
            out.put(']');
            break;
        case Tv::Struct:
            out.put('{');
            for (std::size_t i = 0; i < v.items.size(); ++i) {
                if (i) out.put(',');
                emit_esc(out, v.keys[i]);
                out.put(':');
                emit_tv(out, v.items[i]);
            }
            out.put('}');
            break;
        case Tv::Json: out.puts(v.s.c_str()); break;
        case Tv::Bytes: break;  // handled above
    }
    if (v.kind == Tv::Enum && !v.s.empty()) {
        out.puts(",\"name\":");
        emit_esc(out, v.s);
    }
    if (!v.name.empty()) {
        out.puts(",\"n\":");
        emit_esc(out, v.name);
    }
    out.put('}');
}

std::string canonical(const Tv& v) {
    Buf b;
    emit_tv(b, v);
    return b.str();
}

// ---- records ----------------------------------------------------------------------------------------

namespace {

void emit_side(Buf& out, const std::vector<SideEntry>& side) {
    out.put('{');
    for (std::size_t i = 0; i < side.size(); ++i) {
        const SideEntry& e = side[i];
        if (i) out.put(',');
        emit_esc(out, e.name);
        out.puts(":{");
        if (e.before_kind == SideEntry::IsNull) {
            out.puts("\"before\":null,");
        } else if (e.before_kind == SideEntry::Value) {
            out.puts("\"before\":");
            emit_tv(out, e.before);
            out.put(',');
        }
        out.puts("\"after\":");
        emit_tv(out, e.after);
        out.put('}');
    }
    out.put('}');
}

void emit_ours_side(Buf& out, const std::vector<OursSide>& side) {
    out.put('{');
    for (std::size_t i = 0; i < side.size(); ++i) {
        if (i) out.put(',');
        emit_esc(out, side[i].name);
        out.puts(":{\"after\":");
        emit_tv(out, side[i].after);
        out.put('}');
    }
    out.put('}');
}

void emit_opt_tv(Buf& out, const std::optional<Tv>& v) {
    if (v) emit_tv(out, *v);
    else out.puts("null");
}

void emit_diff(Buf& out, const std::vector<DiffEntry>& diff) {
    out.put('[');
    for (std::size_t i = 0; i < diff.size(); ++i) {
        const DiffEntry& d = diff[i];
        if (i) out.put(',');
        out.puts("{\"path\":");
        emit_esc(out, d.path);
        out.puts(",\"why\":");
        emit_esc(out, d.why, std::strlen(d.why));
        out.puts(",\"orig\":");
        if (!d.orig_raw.empty()) out.puts(d.orig_raw.c_str());
        else emit_opt_tv(out, d.orig);
        out.puts(",\"ours\":");
        if (!d.ours_raw.empty()) out.puts(d.ours_raw.c_str());
        else emit_opt_tv(out, d.ours);
        if (d.first_diff_offset >= 0) out.printf(",\"first_diff_offset\":%lld", d.first_diff_offset);
        if (!d.note.empty()) {
            out.puts(",\"note\":");
            emit_esc(out, d.note);
        }
        out.put('}');
    }
    out.put(']');
}

}  // namespace

// mkfixture.emit_record(): fixed KEY_ORDER, optional keys omitted.
void emit_record(Buf& out, const Record& r) {
    out.printf("{\"ts\":%lld,\"hook\":", static_cast<long long>(r.ts));
    emit_esc(out, r.hook);
    out.printf(",\"mode\":\"%s\",\"call_id\":%llu,\"thread\":%llu", mode_name(r.mode),
               static_cast<unsigned long long>(r.call_id), static_cast<unsigned long long>(r.thread));
    if (r.depth >= 0) out.printf(",\"depth\":%d", r.depth);
    out.puts(",\"args\":[");
    for (std::size_t i = 0; i < r.args.size(); ++i) {
        if (i) out.put(',');
        emit_tv(out, r.args[i]);
    }
    out.puts("],\"ret\":");
    emit_opt_tv(out, r.ret);
    out.puts(",\"side\":");
    emit_side(out, r.side);
    if (r.has_ours) {
        out.puts(",\"ours\":{\"ret\":");
        emit_opt_tv(out, r.ours_ret);
        out.puts(",\"side\":");
        emit_ours_side(out, r.ours_side);
        out.put('}');
    }
    if (r.diverged >= 0) out.puts(r.diverged ? ",\"diverged\":true" : ",\"diverged\":false");
    if (r.has_diff) {
        out.puts(",\"diff\":");
        emit_diff(out, r.diff);
    }
    if (r.err) {
        out.puts(",\"err\":");
        emit_esc(out, *r.err);
    }
    if (r.note) {
        out.puts(",\"note\":");
        emit_esc(out, *r.note);
    }
    out.puts("}\n");
}

// mkfixture.emit_meta() over the dict shape mkfixture.meta() builds.
void emit_meta(Buf& out, const Meta& m) {
    out.puts("{\"meta\":{\"format\":1,\"build\":");
    emit_esc(out, m.build);
    out.puts(",\"exe_sha256\":");
    emit_esc(out, m.exe_sha256);
    out.puts(",\"started\":");
    emit_esc(out, m.started);
    out.printf(",\"inline_max\":%u,\"hooks\":{", m.inline_max);
    for (std::size_t i = 0; i < m.hooks.size(); ++i) {
        const HookPolicy& p = m.hooks[i].second;
        if (i) out.put(',');
        emit_esc(out, m.hooks[i].first);
        out.puts(":{\"ftol\":");
        emit_float(out, p.ftol, false);
        out.printf(",\"ftol_kind\":\"%s\",\"ptr\":\"%s\"", p.ftol_kind, p.ptr_exact ? "exact" : "ignore");
        if (!p.unordered.empty()) {
            out.puts(",\"unordered\":[");
            for (std::size_t j = 0; j < p.unordered.size(); ++j) {
                if (j) out.put(',');
                emit_esc(out, p.unordered[j]);
            }
            out.put(']');
        }
        out.put('}');
    }
    out.puts("}}}\n");
}

}  // namespace shim::trace