#!/usr/bin/env python3
"""save_writer_stub.py -- synthetic SOTS1 save builder for testing save_reader.

This is NOT a game-compatible writer.  It emits the same *framing* the reader
assumes (name-tagged values, NUL padding, BEEFBEEF/41104110 frames, gzip) and
drives itself off the reader's schema so every declared struct gets exercised.
Values are deterministic counters, so the expected typed dict is known up
front and the round trip can be asserted.  Real-save validation is pending.

    python3 save_writer_stub.py out.sav [--padding joint|split] [--inflated]

Stdlib only.
"""
from __future__ import annotations

import argparse
import gzip
import random
import struct
import sys

import save_reader as sr

__all__ = ["SaveWriter", "Fixture", "build_fixture"]


class SaveWriter:
    """Byte builder for the Streamable framing.  padding = 'joint' | 'split'."""

    def __init__(self, padding: str = "joint"):
        if padding not in ("joint", "split"):
            raise ValueError(padding)
        self.padding = padding
        self.buf = bytearray()
        self.depth = 0

    # -- primitives ----------------------------------------------------------------
    def _tag(self, name: str) -> int:
        b = name.encode("ascii")
        self.buf += struct.pack("<i", len(b)) + b
        return len(b)

    def _pad(self):
        self.buf += b"\0" * (sr.pad4(len(self.buf)) - len(self.buf))

    def item(self, name: str, payload: bytes) -> int:
        """Named value.  Returns the offset of the value bytes."""
        start = len(self.buf)
        self._tag(name)
        if self.padding == "split":
            self._pad()
        vp = len(self.buf)
        self.buf += payload
        self._pad()                        # joint: pad over tag+payload; split: pad payload
        assert len(self.buf) == sr.pad4(len(self.buf)) and start % 4 == 0
        return vp

    def int(self, name, v):
        return self.item(name, struct.pack("<i", v))

    def uint(self, name, v):
        return self.item(name, struct.pack("<I", v))

    def float(self, name, v):
        return self.item(name, struct.pack("<f", v))

    def bool(self, name, v):
        return self.item(name, b"\1" if v else b"\0")

    def int64(self, name, v):
        return self.item(name, struct.pack("<q", v))

    def string(self, name, s):
        b = s.encode("cp1252")
        return self.item(name, struct.pack("<i", len(b)) + b)

    def raw(self, name, data: bytes):
        return self.item(name, data)

    def begin(self, name):
        """Open a frame (name=None -> tagless frame)."""
        if name is not None:
            self._tag(name)
            self._pad()
        self.buf += sr.BEGIN_BYTES
        self.depth += 1

    def end(self):
        assert self.depth > 0
        self.buf += sr.END_BYTES
        self.depth -= 1

    def vec3(self, name, x, y, z, named: bool = False):
        self.begin(name)
        if named:
            self.float(".", x)
            self.float(".", y)
            self.float(".", z)
        else:
            self.buf += struct.pack("<3f", x, y, z)
        self.end()

    def bytes(self) -> bytes:
        assert self.depth == 0, "unbalanced frames"
        return bytes(self.buf)

    def gzip(self) -> bytes:
        return gzip.compress(self.bytes(), mtime=0)


# --- schema-driven fixture ----------------------------------------------------

GATES_TRUE = {"vnh", "hindi", "isind", "HFPlan", "HLay", "hbq", "hsp", "HasAIR"}
OPT_EMIT = {"BQ"}                  # optional items the fixture does write
ARRAY_LEN = 2


def _plain_scalar(name, kind, value):
    return {"name": name, "kind": kind, "value": value}


class Fixture:
    """Emits a synthetic save from sr.ROOT and records the expected typed dict."""

    def __init__(self, padding="joint", seed=1):
        self.w = SaveWriter(padding)
        self.counter = 0
        self.rng = random.Random(seed)
        self.expected = None

    # -- value policy ----------------------------------------------------------------
    def next(self) -> int:
        self.counter += 1
        return self.counter

    def value(self, name, kind, ctx):
        c = self.next()
        if kind == "int":
            if name == "idx" and ctx is sr.PlayerColor:
                self.color_toggle = not getattr(self, "color_toggle", False)
                return -1 if self.color_toggle else 3      # alternate custom-RGB / palette
            return c if c % 7 else -c
        if kind == "float":
            return c + 0.25
        if kind == "bool":
            return True if name in GATES_TRUE else bool(c % 2)
        if kind == "int64":
            return c * (1 << 33) + 7
        if kind == "string":
            return "str%d%s" % (c, "x" * (c % 4)) + ("é" if c % 5 == 0 else "")
        raise ValueError(kind)

    def emit_scalar(self, name, kind, ctx):
        v = self.value(name, kind, ctx)
        getattr(self.w, kind)(name, v)
        return v

    # -- generic filler ----------------------------------------------------------------
    def generic_frame(self, name, with_blob=False):
        """Frame with content the schema knows nothing about.  Returns plain()."""
        w = self.w
        w.begin(name)
        items = []
        if with_blob:
            blob = bytes(self.rng.getrandbits(8) for _ in range(2500))
            while sr.END_BYTES in blob or sr.BEGIN_BYTES in blob:
                blob = bytes(self.rng.getrandbits(8) for _ in range(2500))
            vp = w.raw("State", blob)
            end = len(w.buf)
            items.append(_plain_scalar("State", "raw", {"len": end - vp, "hex": w.buf[vp:vp + 32].hex()}))
        else:
            c = self.next()
            w.int("ga", c)
            items.append(_plain_scalar("ga", "int", c))
            w.bool("cmp", True)                         # 'cmp' is in the catalog -> bool
            items.append(_plain_scalar("cmp", "bool", True))
            w.string("gs", "gen%d" % c)
            items.append(_plain_scalar("gs", "string", "gen%d" % c))
            w.begin("gn")
            w.float("gflt", c + 0.5)
            w.end()
            items.append({"_name": "gn", "_items": [_plain_scalar("gflt", "float", c + 0.5)]})
            w.begin(None)
            w.int("ti", 42)
            w.end()
            items.append({"_name": None, "_items": [_plain_scalar("ti", "int", 42)]})
        w.end()
        return {"_name": name, "_items": items}

    def filler_items(self):
        c = self.next()
        self.w.int("Filler", c)
        out = [_plain_scalar("Filler", "int", c)]
        out.append(self.generic_frame("FillerF"))
        return out

    # -- schema walk -------------------------------------------------------------------
    def build(self) -> bytes:
        self.expected = {}
        self.emit_fields(sr.ROOT.fields, self.expected, None, top=True)
        return self.w.bytes()

    def emit_fields(self, fields, out: dict, ctx, top=False):
        for f in fields:
            if isinstance(f, sr.Field):
                self.emit_field(f, out, ctx)
            elif isinstance(f, sr.Opt):
                if f.field.name in OPT_EMIT:
                    self.emit_field(f.field, out, ctx)
            elif isinstance(f, sr.If):
                if f.holds(out):
                    inner = f.inner
                    if isinstance(inner, sr.Seq):
                        self.emit_fields(inner.fields, out, ctx)
                    else:
                        self.emit_field(inner, out, ctx)
            elif isinstance(f, sr.Repeat):
                out[f.key] = [self.emit_elem(f.elem, ctx) for _ in range(ARRAY_LEN)]
            elif isinstance(f, sr.Until):
                out[f.key] = self.filler_items()
            elif isinstance(f, sr.Rest):
                pass                                    # nothing extra is written
            else:
                raise TypeError(f)

    def emit_field(self, f: sr.Field, out: dict, ctx):
        t = f.type
        if isinstance(t, sr.NArr):
            self.w.int(f.name, ARRAY_LEN)
            out[f.key] = [self.emit_elem(t.elem, ctx) for _ in range(ARRAY_LEN)]
        elif isinstance(t, sr.Seq):
            self.emit_fields(t.fields, out, ctx)
        elif isinstance(t, sr.Shape):
            if f.flex and self.counter % 2:
                sub = {}
                self.emit_fields(t.fields, sub, t)          # inline variant
                out[f.key] = sub
            else:
                # authoritative fields use their disk tag; R1-only fields
                # borrow the Shape's (guessed) frame name where it has one
                out[f.key] = self.emit_shape(f.name if f.auth else (t.name or f.name), t)
        elif isinstance(t, sr.CArr):
            if f.flex and self.counter % 2:
                self.w.int(f.name, ARRAY_LEN)                # inline variant == NArr
                out[f.key] = [self.emit_elem(t.elem, ctx) for _ in range(ARRAY_LEN)]
            else:
                out[f.key] = self.emit_carr(f.name, t, ctx)
        elif t == "vec3":
            c = self.next()
            v = [c + 0.5, c + 1.5, c + 2.5]
            self.w.vec3(f.name, *v, named=bool(c % 3 == 0))
            out[f.key] = v
        elif t == "any":
            out[f.key] = self.generic_frame(f.name, with_blob=(f.name in sr.RAW_FRAMES))
        elif t in sr.PRIMITIVES:
            out[f.key] = self.emit_scalar(f.name, t, ctx)
        else:
            raise TypeError(t)

    def emit_shape(self, name, shape: sr.Shape) -> dict:
        self.w.begin(name)
        sub = {}
        self.emit_fields(shape.fields, sub, shape)
        self.w.end()
        return sub

    def emit_carr(self, name, carr: sr.CArr, ctx) -> list:
        self.w.begin(name)
        self.w.int(".", ARRAY_LEN)                        # VectorHelper count tag is "."
        items = [self.emit_elem(carr.elem, ctx) for _ in range(ARRAY_LEN)]
        self.w.end()
        return items

    def emit_elem(self, elem, ctx):
        if isinstance(elem, sr.Seq):
            sub = {}
            self.emit_fields(elem.fields, sub, ctx)
            return sub
        if isinstance(elem, sr.Field):
            sub = {}
            self.emit_field(elem, sub, ctx)
            return sub.get(elem.key)
        if isinstance(elem, sr.Shape):
            return self.emit_shape(".", elem)             # VectorHelper element tag is "."
        if isinstance(elem, sr.CArr):
            return self.emit_carr(".", elem, ctx)
        if elem == "any":
            return self.generic_frame(".")
        if elem in sr.PRIMITIVES:
            return self.emit_scalar(".", elem, ctx)
        raise TypeError(elem)


def build_fixture(padding="joint", seed=1):
    """-> (inflated bytes, expected typed dict)"""
    fx = Fixture(padding, seed)
    data = fx.build()
    return data, fx.expected


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="write a synthetic SOTS1 save for reader tests")
    ap.add_argument("out")
    ap.add_argument("--padding", choices=("joint", "split"), default="joint")
    ap.add_argument("--inflated", action="store_true", help="write the raw stream, not gzip")
    args = ap.parse_args(argv)
    data, _ = build_fixture(args.padding)
    with open(args.out, "wb") as f:
        f.write(data if args.inflated else gzip.compress(data, mtime=0))
    print(f"wrote {args.out}: {len(data)} inflated bytes, padding={args.padding}")
    return 0


if __name__ == "__main__":
    sys.exit(main())