"""paralogue_decoder.py - shared protocol-decoding primitives.

This module concentrates the framing, PacketHead parsing, and msg_id ->
protobuf-type identification used by both the Frida in-process decoder
(run_decoder.py) and the MITM proxy (tcp_proxy_decrypt.py).

Wire format (confirmed against live plaintext):
    [2B BE prefix N][N-byte PacketHead protobuf][body_len bytes body]

`PacketHead.body_len` is body-only (not total).
"""

from __future__ import annotations

import json
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
if SCRIPT_DIR not in sys.path:
    sys.path.insert(0, SCRIPT_DIR)

import paralogue_net_full_pb2 as pb_mod  # noqa: E402
from google.protobuf.message import DecodeError  # noqa: E402

KNOWN_IDS_PATH = os.path.join(SCRIPT_DIR, "known_msg_ids.json")
PACKET_HEAD_CLS = pb_mod.PacketHead


def build_catalog():
    fd = pb_mod.DESCRIPTOR
    out = []
    for name, desc in fd.message_types_by_name.items():
        cls = getattr(pb_mod, name, None)
        if cls is None:
            continue
        out.append((name, desc, cls))
    return out


CATALOG = build_catalog()


def load_known_ids() -> dict[int, str]:
    if not os.path.isfile(KNOWN_IDS_PATH):
        return {}
    try:
        with open(KNOWN_IDS_PATH, "r", encoding="utf-8") as f:
            return {int(k): v for k, v in json.load(f).items()}
    except Exception:
        return {}


def save_known_ids(m: dict[int, str]) -> None:
    try:
        with open(KNOWN_IDS_PATH, "w", encoding="utf-8") as f:
            json.dump({str(k): v for k, v in sorted(m.items())}, f, indent=2)
    except Exception:
        pass


# ---------------------------------------------------------------------------
# msg_id identification
# ---------------------------------------------------------------------------

_TYPE_STRING = 9
_NUMERIC_TYPES = {1, 2, 3, 4, 5, 6, 7, 13, 14, 15, 16, 17, 18}
_LABEL_REPEATED = 3


def _score_fit(msg, body: bytes) -> int:
    try:
        set_fields = msg.ListFields()
    except Exception:
        return 0
    if not set_fields:
        return 0

    score = len(set_fields) * 2
    try:
        rt = msg.SerializeToString()
    except Exception:
        rt = b""
    if rt == body:
        score += 100
    elif rt and len(rt) == len(body):
        score += 20
    elif rt:
        score += int(min(len(rt), len(body)) / max(len(rt), len(body)) * 10)

    if rt:
        missing = max(0, len(body) - len(rt))
        score -= missing // 4

    for f, v in set_fields:
        if f.type == _TYPE_STRING:
            vals = v if f.label == _LABEL_REPEATED else [v]
            for s in vals:
                if not isinstance(s, str):
                    continue
                try:
                    s.encode("utf-8")
                except Exception:
                    score -= 50
                    continue
                if s and sum(1 for c in s if c.isprintable()) / len(s) >= 0.8:
                    score += 5
                elif s and sum(1 for c in s if c.isprintable()) / len(s) < 0.5:
                    score -= 10

        if (f.label == _LABEL_REPEATED
                and f.type in _NUMERIC_TYPES
                and isinstance(v, (list, tuple)) and v):
            try:
                ints = [int(x) for x in v]
            except Exception:
                ints = []
            if ints and all(32 <= x <= 126 for x in ints) and len(ints) >= 3:
                score -= 40

    return score


def identify_body(body: bytes, preferred_name: Optional[str] = None):
    """Return (name, parsed, score, candidates, top3)."""
    if not body:
        return None, None, 0, 0, []

    scored = []
    for name, _desc, cls in CATALOG:
        try:
            m = cls()
            m.MergeFromString(body)
        except (DecodeError, ValueError, Exception):
            continue
        s = _score_fit(m, body)
        if s > 0:
            scored.append((s, name, m))
    if not scored:
        return None, None, 0, 0, []

    scored.sort(key=lambda t: (-t[0], 0 if t[1] == preferred_name else 1, t[1]))
    if preferred_name:
        for s, n, m in scored:
            if n == preferred_name:
                top = [(nn, ss) for ss, nn, _ in scored[:3]]
                return n, m, s, len(scored), top

    best_score, best_name, best_msg = scored[0]
    top = [(nn, ss) for ss, nn, _ in scored[:3]]
    return best_name, best_msg, best_score, len(scored), top


def _has_packed_ascii_string(msg) -> bool:
    """True if a repeated-numeric field contains a run of printable
    ASCII bytes long enough to look like a string. Misidentified string
    fields decode as `repeated int32` where the varints are the UTF-8
    byte values, preceded by tag+length bytes (often small non-printable
    numbers like 10, 7). We scan for any consecutive run of >= 5
    printable ASCII values anywhere in the list.

    Note: protobuf-python returns `RepeatedScalarFieldContainer` for
    repeated fields; it's iterable and supports len() but is NOT a list
    or tuple, so don't use isinstance for the type check.
    """
    try:
        for f, v in msg.ListFields():
            if f.label != _LABEL_REPEATED or f.type not in _NUMERIC_TYPES:
                continue
            try:
                n = len(v)
            except TypeError:
                continue
            if n < 5:
                continue
            try:
                ints = [int(x) for x in v]
            except Exception:
                continue
            run = 0
            for x in ints:
                if 32 <= x <= 126:
                    run += 1
                    if run >= 5:
                        return True
                else:
                    run = 0
    except Exception:
        pass
    return False


def identify_body_strict(body: bytes):
    """Strict identification: return list of candidate type names whose
    protobuf parse consumes the entire body with zero unknown fields.

    A type T is strict-valid for body B if:
      * T().MergeFromString(B) succeeds, AND
      * SerializeToString() before DiscardUnknownFields equals after.

    Empty bodies return []. Types that parse but leave unknown fields
    are excluded: they "absorbed" bytes they can't structurally account
    for, which means they're the wrong type.
    """
    if not body:
        return []
    out = []
    for name, _desc, cls in CATALOG:
        try:
            m = cls()
            m.MergeFromString(body)
        except (DecodeError, ValueError, Exception):
            continue
        try:
            before = m.SerializeToString()
            m.DiscardUnknownFields()
            after = m.SerializeToString()
        except Exception:
            continue
        if before != after:
            continue  # had unknown fields -> wrong type
        if _has_packed_ascii_string(m):
            continue  # misread string as packed-int repeated -> wrong type
        # Demand the serialized form is byte-identical to the input.
        # Protobuf fields are unordered on the wire but generated code
        # emits in a stable order; for the correct type, MergeFromString
        # + SerializeToString round-trips to the canonical form of the
        # same content. We accept equal-LENGTH as a weaker match because
        # the server may have emitted fields in a different order.
        if before == body or len(before) == len(body):
            out.append((name, m, before == body))
    return out


def msg_to_dict(msg) -> dict:
    from google.protobuf import json_format
    try:
        return json_format.MessageToDict(
            msg,
            preserving_proto_field_name=True,
            including_default_value_fields=False,
        )
    except Exception:
        out = {}
        for f, v in msg.ListFields():
            out[f.name] = v.hex() if isinstance(v, bytes) else repr(v)
        return out


# ---------------------------------------------------------------------------
# stream reassembly
# ---------------------------------------------------------------------------

@dataclass
class Stream:
    """One reassembled byte stream."""
    label: str                    # e.g. "C->S" or "S->C"
    chunks: dict = field(default_factory=dict)
    consumed: int = 0
    frontier: int = 0
    _cache_bytes: bytearray = field(default_factory=bytearray)

    def add(self, off: int, data: bytes) -> None:
        if not self.chunks:
            self.consumed = off
            self.frontier = off
        end = off + len(data)
        if end > len(self._cache_bytes):
            self._cache_bytes.extend(b"\x00" * (end - len(self._cache_bytes)))
        self._cache_bytes[off:end] = data
        self.chunks[off] = len(data)
        f = self.frontier
        while True:
            advanced = False
            for co, cl in self.chunks.items():
                if co <= f < co + cl:
                    f = co + cl
                    advanced = True
                    break
            if not advanced:
                break
        self.frontier = f

    def available(self) -> int:
        return self.frontier - self.consumed

    def peek(self, n: int) -> bytes:
        return bytes(self._cache_bytes[self.consumed:self.consumed + n])

    def consume(self, n: int) -> bytes:
        out = bytes(self._cache_bytes[self.consumed:self.consumed + n])
        self.consumed += n
        return out


def drain_stream(stream: Stream):
    """Yield (head, body, skipped) for each complete message.

    `head` is a parsed PacketHead, `body` is bytes, `skipped` is the number
    of junk bytes the framing recovered from before this packet (usually 0).
    """
    while True:
        avail = stream.available()
        if avail < 2:
            return
        prefix = stream.peek(2)
        hlen = (prefix[0] << 8) | prefix[1]
        if hlen == 0 or hlen > 32:
            stream.consume(1)
            continue
        if avail < 2 + hlen:
            return
        head_bytes = bytes(
            stream._cache_bytes[stream.consumed + 2:
                                stream.consumed + 2 + hlen]
        )
        head = PACKET_HEAD_CLS()
        try:
            head.MergeFromString(head_bytes)
        except DecodeError:
            stream.consume(1)
            continue
        try:
            rt = head.SerializeToString()
        except Exception:
            stream.consume(1)
            continue
        if len(rt) != hlen:
            stream.consume(1)
            continue
        mid = int(head.msg_id)
        blen = int(head.body_len)
        if mid <= 0 or mid > 100000 or blen < 0 or blen > (1 << 20):
            stream.consume(1)
            continue
        if avail < 2 + hlen + blen:
            return

        stream.consume(2 + hlen)
        body = stream.consume(blen)
        yield head, body, 0


def short_repr(d) -> str:
    if isinstance(d, dict):
        parts = []
        for k, v in list(d.items())[:4]:
            if isinstance(v, (dict, list)):
                parts.append(f"{k}=<{type(v).__name__}>")
            else:
                s = str(v)
                if len(s) > 40:
                    s = s[:37] + "..."
                parts.append(f"{k}={s}")
        extra = "" if len(d) <= 4 else f" (+{len(d)-4} more)"
        return "{" + ", ".join(parts) + "}" + extra
    return str(d)