#include "reader.h"

#include <cmath>
#include <cstring>

namespace mars::stream {

namespace {

bool cp1252_defined(uint8_t c) {
    return !(c == 0x81 || c == 0x8d || c == 0x8f || c == 0x90 || c == 0x9d);
}

// Guard used only while GUESSING that an unknown tag holds a string: every
// byte windows-1252 defines counts as text (system names carry 0x92).
bool text_plausible(const uint8_t* b, size_t n) {
    if (n == 0) return true;
    size_t ok = 0;
    for (size_t i = 0; i < n; ++i) {
        uint8_t c = b[i];
        if ((c >= 0x20 && c < 0x7f) || (c >= 0x80 && cp1252_defined(c)) || c == 9 || c == 10 || c == 13) ++ok;
    }
    return ok * 10 >= n * 9;
}

// Unknown 4-byte word -> int or float by bit pattern.
Prim classify_word(const uint8_t* raw) {
    int32_t i = rd_i32(raw);
    double f = rd_f32(raw);
    if (i >= -kWordIntAbs && i <= kWordIntAbs) return Prim::Int;
    if (std::isfinite(f) && std::fabs(f) >= 1e-6 && std::fabs(f) < 1e12) return Prim::Float;
    return Prim::Int;
}

Kind prim_to_kind(Prim p) {
    switch (p) {
        case Prim::Int: return Kind::Int;
        case Prim::Float: return Kind::Float;
        case Prim::Bool: return Kind::Bool;
        case Prim::Int64: return Kind::Int64;
        case Prim::String: return Kind::String;
        case Prim::Raw: return Kind::Raw;
        default: return Kind::Int;
    }
}

std::string lower(std::string s) {
    for (char& c : s)
        if (c >= 'A' && c <= 'Z') c = char(c - 'A' + 'a');
    return s;
}

}  // namespace

Walker::Walker(const uint8_t* data, size_t n, const Registry* reg)
    : d_(data), n_(static_cast<uint32_t>(n)), reg_(reg) {}

void Walker::issue(Issue::Level lvl, uint32_t off, std::string msg, const std::string& path) {
    issues.push_back(Issue{lvl, path, off, std::move(msg)});
}

bool Walker::zero(uint32_t a, uint32_t b) const {
    for (uint32_t i = a; i < b; ++i)
        if (d_[i]) return false;
    return true;
}

Walker::TagAt Walker::tag_at(uint32_t p, bool allow_empty) const {
    TagAt t;
    if (p + 4 > n_) return t;
    int32_t ln = rd_i32(d_ + p);
    if (ln == 0) {
        if (!allow_empty) return t;
        t.ok = true;
        t.end = p + 4;
        return t;
    }
    if (ln < 0 || ln > kMaxTagLen || uint64_t(p) + 4 + uint32_t(ln) > n_) return t;
    const uint8_t* b = d_ + p + 4;
    for (int i = 0; i < ln; ++i)
        if (b[i] < 0x20 || b[i] >= 0x7f) return t;
    t.ok = true;
    t.name.assign(reinterpret_cast<const char*>(b), size_t(ln));
    t.end = p + 4 + uint32_t(ln);
    return t;
}

// Read a candidate scalar at the value position (joint padding: right after the
// name).  known=true means the kind comes from the schema: a string value is
// then accepted whatever bytes it holds (only tags must be ASCII).
Walker::Scalar Walker::try_scalar(uint32_t tag_start, uint32_t name_end, Cand kind, bool known) {
    Scalar r;
    uint32_t vp = name_end;
    uint32_t size;
    if (kind == Cand::String) {
        if (vp + 4 > n_) {
            eof_limited_ = true;
            return r;
        }
        int32_t ln = rd_i32(d_ + vp);
        if (ln < 0 || ln > kMaxStringLen) return r;
        if (uint64_t(vp) + 4 + uint32_t(ln) > n_) {
            eof_limited_ = true;
            return r;
        }
        if (!known && !text_plausible(d_ + vp + 4, size_t(ln))) return r;
        size = 4 + uint32_t(ln);
    } else {
        size = kind == Cand::Word ? 4 : kind == Cand::Bool ? 1 : 8;
        if (uint64_t(vp) + size > n_) {
            eof_limited_ = true;
            return r;
        }
        if (kind == Cand::Bool && d_[vp] > 1) return r;
    }
    uint32_t end = tag_start + pad4(name_end - tag_start + size);
    if (end > n_) {
        eof_limited_ = true;
        return r;
    }
    if (!zero(vp + size, end)) return r;
    r.ok = true;
    r.cand = kind;
    r.vp = d_ + vp;
    r.vsize = size;
    r.end = end;
    return r;
}

// Does an item plausibly start at p?  Looks `depth` items ahead.
bool Walker::plausible_item(uint32_t p, int depth, bool allow_empty) {
    if (p == n_) return true;
    if (p + 4 > n_) {
        eof_limited_ = true;
        return false;
    }
    uint32_t w = u32(p);
    if (w == kEndMark || w == kBeginMark) return true;
    TagAt t = tag_at(p, allow_empty);
    if (!t.ok) return false;
    uint32_t a = pad4(t.end);
    if (a + 4 <= n_ && u32(a) == kBeginMark && zero(t.end, a)) return true;
    if (depth <= 0) return true;
    // a tag the catalog knows to be a string is read as one without any text
    // test, so a non-ASCII value can never veto the layout of the item before it
    if (reg_ && reg_->kind(t.name) == Prim::String) {
        Scalar r = try_scalar(p, t.end, Cand::String, true);
        if (r.ok && plausible_item(r.end, depth - 1, allow_empty)) return true;
    }
    static const Cand order[] = {Cand::Word, Cand::Bool, Cand::String, Cand::Int64};
    for (Cand c : order) {
        Scalar r = try_scalar(p, t.end, c, false);
        if (!r.ok) continue;
        if (plausible_item(r.end, depth - 1, allow_empty)) return true;
    }
    return false;
}

// Scan forward for the next END marker or plausible tag.
std::pair<const char*, uint32_t> Walker::resync(uint32_t p, int depth) {
    uint32_t q = p + 1;
    while (q + 4 <= n_) {
        uint32_t w = u32(q);
        if (w == kEndMark && depth > 0) return {"end", q};
        if (w == kBeginMark) return {"begin", q};
        if (tag_at(q, false).ok && plausible_item(q, 2, false)) return {"tag", q};
        ++q;
    }
    return {"eof", n_};
}

Hint Walker::child_hint(const Desc* frame, size_t index, const std::string* name) const {
    Hint h;
    if (frame) {
        if (frame->type == Desc::Shape) {
            if (index < frame->prefix.size()) h = frame->prefix[index];
            if (h.empty() && name && !name->empty()) {
                auto it = frame->by_name.find(*name);
                if (it == frame->by_name.end()) it = frame->by_name.find(lower(*name));
                if (it != frame->by_name.end()) h = it->second;
            }
        } else if (frame->type == Desc::CArr) {
            if (index == 0) h.kind = Prim::Int;
            else h = frame->elem;
        } else if (frame->type == Desc::Raw) {
            h.kind = Prim::Raw;
        }
    }
    if (h.empty() && name && !name->empty() && reg_) {
        // global catalog: exact case only (case-folding was found to mislabel
        // unknown names, e.g. 'a' vs star-colour 'A')
        h.kind = reg_->kind(*name);
        h.sub = reg_->shape(*name);
    }
    return h;
}

Node Walker::walk() {
    uint32_t p = 0;
    Node root;
    root.tagged = false;
    root.kind = Kind::Complex;
    root.children = walk_frame(p, false, "", 0, nullptr, "");
    root.size = p;
    return root;
}

// Read items until the frame's END marker (or EOF).  p is left after the END.
std::vector<Node> Walker::walk_frame(uint32_t& p, bool has_ctx, const std::string& ctx, int depth,
                                     const Desc* shape, const std::string& path) {
    std::vector<Node> children;
    size_t index = 0;
    if (depth > 0) ++stats.frames;
    std::string ctxr = has_ctx ? "'" + ctx + "'" : "None";
    for (;;) {
        if (p >= n_) {
            if (depth > 0) issue(Issue::Error, p, "frame " + ctxr + " not terminated before EOF", path);
            return children;
        }
        if (p + 4 > n_) {
            children.push_back(raw_node(false, "", p, n_, "trailing bytes", Issue::Warn, path));
            if (depth > 0) issue(Issue::Error, n_, "frame " + ctxr + " not terminated before EOF", path);
            p = n_;
            return children;
        }
        uint32_t w = u32(p);
        if (w == kEndMark) {
            if (depth > 0) {
                p += 4;
                return children;
            }
            issue(Issue::Warn, p, "stray END marker at top level", path);
            children.push_back(raw_node(false, "", p, p + 4, "stray END marker", Issue::Warn, path));
            p += 4;
            continue;
        }
        children.push_back(read_item(p, index, depth, shape, path));
        ++index;
    }
}

Node Walker::read_item(uint32_t& p, size_t index, int depth, const Desc* frame, const std::string& path) {
    ++stats.items;
    uint32_t start = p;
    uint32_t w = u32(p);
    if (w == kBeginMark) {  // tagless frame
        Hint h = child_hint(frame, index, nullptr);
        Node n;
        n.tagged = false;
        n.kind = Kind::Complex;
        n.offset = start;
        p += 4;
        n.children = walk_frame(p, false, "", depth + 1, h.sub, path + "/<" + std::to_string(index) + ">");
        n.size = p - start;
        return n;
    }
    TagAt t = tag_at(p, true);
    if (!t.ok) return unreadable(p, depth, path);
    const std::string& name = t.name;
    uint32_t ne = t.end;
    Hint h = child_hint(frame, index, &name);
    uint32_t a = pad4(ne);
    if (a + 4 <= n_ && u32(a) == kBeginMark && zero(ne, a)) {
        Node n;
        n.name = name;
        n.kind = Kind::Complex;
        n.offset = start;
        n.hinted = h.sub != nullptr;
        p = a + 4;
        n.children = walk_frame(p, true, name, depth + 1, h.sub, path + "/" + name);
        n.size = p - start;
        return n;
    }

    // scalar: hinted kind first, then the guess order
    if (h.kind == Prim::Raw) {
        uint32_t vp = ne;
        uint32_t q = n_;
        if (depth > 0) {
            // the blob runs to the next END marker (any byte alignment, like bytes.find)
            const uint8_t endb[4] = {0x10, 0x41, 0x10, 0x41};
            for (uint32_t i = vp; i + 4 <= n_; ++i)
                if (std::memcmp(d_ + i, endb, 4) == 0) {
                    q = i;
                    break;
                }
        }
        Node n;
        n.name = name;
        n.kind = Kind::Raw;
        n.offset = start;
        n.size = q - start;
        n.hinted = true;
        n.raw.assign(d_ + vp, d_ + q);
        stats.raw_bytes += q - vp;
        p = q;
        return n;
    }

    Cand order[4];
    int norder = 0;
    switch (h.kind) {
        case Prim::Int:
        case Prim::Float: order[norder++] = Cand::Word; break;
        case Prim::Bool: order[norder++] = Cand::Bool; break;
        case Prim::String: order[norder++] = Cand::String; break;
        case Prim::Int64: order[norder++] = Cand::Int64; break;
        default: break;
    }
    for (Cand c : {Cand::Word, Cand::Bool, Cand::String, Cand::Int64}) {
        bool have = false;
        for (int i = 0; i < norder; ++i)
            if (order[i] == c) have = true;
        if (!have) order[norder++] = c;
    }
    bool hinted_kind = h.kind != Prim::None;

    // pass 1: continuation must be a non-empty tag / marker; pass 2 tolerates
    // empty ("") tags; pass 3 (only when lookahead was defeated by EOF) takes
    // whatever fits the remaining bytes.
    eof_limited_ = false;
    Scalar chosen;
    int chosen_i = -1;
    int chosen_mode = -1;
    for (int mode = 0; mode < 3 && chosen_i < 0; ++mode) {
        if (mode == 2 && !eof_limited_) break;
        for (int i = 0; i < norder; ++i) {
            Scalar r = try_scalar(start, ne, order[i], i == 0 && hinted_kind);
            if (!r.ok) continue;
            if (mode == 2 || plausible_item(r.end, 2, mode == 1)) {
                chosen = r;
                chosen_i = i;
                chosen_mode = mode;
                break;
            }
        }
    }
    if (chosen_i >= 0) {
        Node n;
        n.name = name;
        n.offset = start;
        n.size = chosen.end - start;
        Prim ck;
        switch (chosen.cand) {
            case Cand::Word:
                ck = (h.kind == Prim::Int || h.kind == Prim::Float) ? h.kind : classify_word(chosen.vp);
                break;
            case Cand::Bool: ck = Prim::Bool; break;
            case Cand::String: ck = Prim::String; break;
            default: ck = Prim::Int64; break;
        }
        n.kind = prim_to_kind(ck);
        if (chosen.cand == Cand::String) n.raw.assign(chosen.vp + 4, chosen.vp + chosen.vsize);
        else n.raw.assign(chosen.vp, chosen.vp + chosen.vsize);
        n.hinted = hinted_kind && chosen_i == 0;
        if (hinted_kind && chosen_i != 0) {
            ++stats.hint_failures;
            issue(Issue::Warn, start,
                  "'" + name + "': hinted type " + prim_name(h.kind) + " not plausible, read as " + prim_name(ck),
                  path);
        } else if (!hinted_kind) {
            ++stats.guessed;
        }
        if (chosen_mode == 2) {
            ++stats.best_effort;
            issue(Issue::Warn, start,
                  "'" + name + "': stream ends before layout can be confirmed; read as " + prim_name(ck), path);
        }
        p = chosen.end;
        return n;
    }

    // tag looked fine but no value layout fits: skip to the next sync point
    auto [what, q] = resync(ne, depth);
    ++stats.resyncs;
    issue(Issue::Warn, start,
          "'" + name + "': no value layout fits; skipped " + std::to_string(q - ne) + " bytes to " + what, path);
    Node n;
    n.name = name;
    n.kind = Kind::Raw;
    n.offset = start;
    n.size = q - start;
    n.raw.assign(d_ + ne, d_ + q);
    stats.raw_bytes += q - ne;
    p = q;
    return n;
}

// No tag at p.  A small unnamed payload before an END is normal (info);
// anything longer is a resync event.
Node Walker::unreadable(uint32_t& p, int depth, const std::string& path) {
    uint32_t start = p;
    if (depth > 0) {
        uint32_t q = start;
        while (q + 4 <= n_ && q - start <= uint32_t(kSmallRaw)) {
            if (u32(q) == kEndMark) {
                p = q;
                return raw_node(false, "", start, q, "unnamed payload", Issue::Info, path);
            }
            q += 4;
        }
    }
    auto [what, q] = resync(start, depth);
    ++stats.resyncs;
    p = q;
    return raw_node(false, "", start, q, (std::string("unreadable; resynced to ") + what).c_str(), Issue::Warn,
                    path);
}

Node Walker::raw_node(bool tagged, const std::string& name, uint32_t p, uint32_t q, const char* why,
                      Issue::Level lvl, const std::string& path) {
    issue(lvl, p, std::string(why) + ": " + std::to_string(q - p) + " raw bytes", path);
    stats.raw_bytes += q - p;
    Node n;
    n.tagged = tagged;
    n.name = name;
    n.kind = Kind::Raw;
    n.offset = p;
    n.size = q - p;
    n.raw.assign(d_ + p, d_ + q);
    return n;
}

Node read_tree(const Bytes& inflated, std::vector<Issue>* issues, Stats* stats, const Registry* reg) {
    Walker w(inflated.data(), inflated.size(), reg);
    Node root = w.walk();
    if (issues) *issues = std::move(w.issues);
    if (stats) *stats = w.stats;
    return root;
}

}  // namespace mars::stream