#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""C08 smoke checks beyond HTTP: does a served ``/predict`` body agree with the stored CPU-reference body of the
same sample, and does ``/info`` report what the staged tt-model package pins for the served profile. Standard
library only.

Like :mod:`client`, this module imports nothing from ttaw, so a bundle's ``smoke_test.py`` (run by
``container_smoke.sh`` with the host's ``python3``, which has neither numpy nor the bundle) loads it by path::

    sys.path.insert(0, str(HERE.parent / "ttaw" / "server"))          # HERE = code/<pkg>/server
    import client as ttaw_client
    import smoke as ttaw_smoke

    info, fails = ttaw_client.check_service(url, expect_dispatch="eth", expect_grid="12x10")
    pinned = ttaw_smoke.pinned_config(f"{staged}/tt_kernel_manifest.json", profile)
    fails += ttaw_smoke.check_pinned(info, pinned, "CENTERPOINT")     # /info runs what serve.env pins
    ref = ttaw_smoke.find_reference(SAMPLE, pinned["profile"], info.get("variant"))
    metrics, more = ttaw_smoke.compare_with_reference(body, ttaw_smoke.load_json(ref), gates={"min_recall": 0.9})

The stored reference is the ``to_dict()`` body of the fp32 CPU reference on the shipped sample, saved next to it as
``<sample stem>.reference.json``, or ``<stem>.<profile>.reference.json`` / ``<stem>.<variant>.reference.json``
when the output depends on the serve profile or the model variant. Synthetic samples may give no detections, so
agreement with the reference, not a detection count, is their smoke gate (PLAN.md section 6.3).

What is compared, by the keys of the reference body (BUNDLE_CONVENTIONS.md section 7.4):

- ``detections``: greedy same-label matching, reference rows in descending score order; 3-D rows (``center``)
  match by BEV centre distance, 2-D rows (``box_xyxy``) by IoU; gates on recall, precision and the largest score
  difference of a matched pair;
- ``trajectory``: average / final displacement over the ``x`` and ``y`` columns;
- every encoded array (``ttaw.io.encode_array`` npz / npz_compressed / raw / list, ``encode_png``), e.g.
  per-point ``labels``, a ``mask`` or a ``semseg`` extra: integer arrays by the fraction of equal elements, float
  arrays by the largest absolute difference;
- ``model`` and ``frame_id`` must be equal.
"""
from __future__ import annotations

import ast
import base64
import io
import json
import math
import re
import struct
import sys
import zipfile
import zlib
from pathlib import Path
from typing import Any, Callable, Dict, List, Mapping, NamedTuple, Optional, Sequence, Tuple, Union

__all__ = ["DEFAULT_GATES", "ARRAY_FORMATS", "Array", "load_json", "decode_array", "find_reference",
           "compare_with_reference", "describe_metrics", "pinned_config", "check_pinned"]

ARRAY_FORMATS = ("npz", "npz_compressed", "raw", "list", "png")
DEFAULT_GATES: Dict[str, Optional[float]] = {
    "max_center_dist": 0.5,       # 3-D detections: same-label BEV centre distance of a match (m)
    "min_iou": 0.5,               # 2-D detections: same-label box IoU of a match
    "min_recall": 0.95,           # matched / reference detections
    "min_precision": 0.95,        # matched / served detections
    "max_score_err": 0.05,        # largest |score difference| of a matched pair
    "min_label_agreement": 0.99,  # integer arrays (per-point labels, masks): fraction of equal elements
    "max_abs_err": None,          # float arrays: largest |difference| (reported only, unless set)
    "max_ade": 0.5,               # trajectories: average displacement over x, y (m)
    "max_fde": 1.0,               # trajectories: final displacement (m)
}

PathLike = Union[str, Path]
# numpy type code (without the byte-order character) -> (struct code, numpy dtype name)
_TYPES = {"b1": ("?", "bool"), "u1": ("B", "uint8"), "u2": ("H", "uint16"), "u4": ("I", "uint32"),
          "u8": ("Q", "uint64"), "i1": ("b", "int8"), "i2": ("h", "int16"), "i4": ("i", "int32"),
          "i8": ("q", "int64"), "f2": ("e", "float16"), "f4": ("f", "float32"), "f8": ("d", "float64")}
_BY_NAME = {name: code for code, name in _TYPES.values()}
_SAFE_KEY = re.compile(r"[A-Za-z0-9_.-]+")


class Array(NamedTuple):
    """A decoded array: ``values`` are the elements in C order (flattened)."""

    shape: Tuple[int, ...]
    dtype: str
    values: List[Any]

    @property
    def is_float(self) -> bool:
        return self.dtype.startswith("float")


def load_json(path: PathLike) -> Dict[str, Any]:
    return json.loads(Path(path).read_text())


# ------------------------------------------------------------------------------------------------ decoding

def _count(shape: Sequence[int]) -> int:
    n = 1
    for d in shape:
        n *= int(d)
    return n


def _fortran_to_c(values: List[Any], shape: Sequence[int]) -> List[Any]:
    """Reorder column-major elements into row-major order."""
    strides, step = [], 1
    for d in shape:
        strides.append(step)
        step *= d
    out, index = [], [0] * len(shape)
    for _ in range(len(values)):
        out.append(values[sum(i * s for i, s in zip(index, strides))])
        for k in range(len(shape) - 1, -1, -1):
            index[k] += 1
            if index[k] < shape[k]:
                break
            index[k] = 0
    return out


def _npy(raw: bytes) -> Array:
    """One ``.npy`` file (format 1.0 / 2.0 / 3.0, plain numeric dtypes)."""
    if raw[:6] != b"\x93NUMPY":
        raise ValueError("not an .npy file")
    major = raw[6]
    if major == 1:
        (hlen,), start = struct.unpack("<H", raw[8:10]), 10
    elif major in (2, 3):
        (hlen,), start = struct.unpack("<I", raw[8:12]), 12
    else:
        raise ValueError(f".npy format version {major} is not supported")
    header = ast.literal_eval(raw[start:start + hlen].decode("utf-8" if major == 3 else "latin1"))
    descr, shape = header["descr"], tuple(int(d) for d in header["shape"])
    if not isinstance(descr, str) or descr[1:] not in _TYPES:
        raise ValueError(f".npy dtype {descr!r} is not supported")
    order = {">": ">", "=": "<" if sys.byteorder == "little" else ">"}.get(descr[0], "<")
    code, name = _TYPES[descr[1:]]
    n = _count(shape)
    data = raw[start + hlen:]
    size = n * struct.calcsize("<" + code)
    if len(data) < size:
        raise ValueError(f".npy data holds {len(data)} bytes, {size} expected")
    values = list(struct.unpack(f"{order}{n}{code}", data[:size]))
    if header.get("fortran_order") and len(shape) > 1:
        values = _fortran_to_c(values, shape)
    return Array(shape, name, values)


def _paeth(a: int, b: int, c: int) -> int:
    p = a + b - c
    pa, pb, pc = abs(p - a), abs(p - b), abs(p - c)
    if pa <= pb and pa <= pc:
        return a
    return b if pb <= pc else c


def _png(raw: bytes) -> Array:
    """An 8-bit, non-interlaced grey / grey+alpha / RGB / RGBA PNG (what ``ttaw.io.encode_png`` writes)."""
    if raw[:8] != b"\x89PNG\r\n\x1a\n":
        raise ValueError("not a PNG file")
    pos, idat, head = 8, [], None
    while pos + 8 <= len(raw):
        length, kind = struct.unpack(">I4s", raw[pos:pos + 8])
        chunk = raw[pos + 8:pos + 8 + length]
        pos += 12 + length
        if kind == b"IHDR":
            head = struct.unpack(">IIBBBBB", chunk)
        elif kind == b"IDAT":
            idat.append(chunk)
        elif kind == b"IEND":
            break
    if head is None:
        raise ValueError("PNG without IHDR")
    width, height, depth, color, _, _, interlace = head
    channels = {0: 1, 2: 3, 4: 2, 6: 4}.get(color)
    if depth != 8 or channels is None or interlace:
        raise ValueError(f"PNG bit depth {depth}, colour type {color}, interlace {interlace} is not supported")
    data = zlib.decompress(b"".join(idat))
    stride, bpp = width * channels, channels
    if len(data) < height * (stride + 1):
        raise ValueError("truncated PNG image data")
    out, prev = bytearray(), bytearray(stride)
    for y in range(height):
        base = y * (stride + 1)
        ftype, row = data[base], bytearray(data[base + 1:base + 1 + stride])
        if ftype == 1:
            for i in range(bpp, stride):
                row[i] = (row[i] + row[i - bpp]) & 0xFF
        elif ftype == 2:
            for i in range(stride):
                row[i] = (row[i] + prev[i]) & 0xFF
        elif ftype == 3:
            for i in range(stride):
                left = row[i - bpp] if i >= bpp else 0
                row[i] = (row[i] + ((left + prev[i]) >> 1)) & 0xFF
        elif ftype == 4:
            for i in range(stride):
                left, upper_left = (row[i - bpp], prev[i - bpp]) if i >= bpp else (0, 0)
                row[i] = (row[i] + _paeth(left, prev[i], upper_left)) & 0xFF
        elif ftype != 0:
            raise ValueError(f"PNG filter type {ftype} is invalid")
        out += row
        prev = row
    shape = (height, width) if channels == 1 else (height, width, channels)
    return Array(shape, "uint8", list(out))


def _flatten(value: Any, out: List[Any]) -> List[Any]:
    if isinstance(value, list):
        for v in value:
            _flatten(v, out)
    else:
        out.append(value)
    return out


def decode_array(spec: Mapping[str, Any]) -> Array:
    """Decode an encoded output array ``{"format", "key", "dtype", "shape", "data"}`` (``ttaw.io.encode_array`` /
    ``encode_png``) without numpy. Raises ``ValueError`` for anything it cannot decode."""
    fmt = spec.get("format")
    if fmt not in ARRAY_FORMATS:
        raise ValueError(f"array format {fmt!r} is not one of {ARRAY_FORMATS}")
    if fmt == "list":
        dtype = str(spec.get("dtype") or "float64")
        arr = Array(tuple(int(d) for d in spec.get("shape") or ()), dtype, _flatten(spec.get("data"), []))
    else:
        raw = base64.b64decode(spec.get("data") or "")
        if fmt == "png":
            arr = _png(raw)
        elif fmt == "raw":
            dtype = str(spec.get("dtype"))
            if dtype not in _BY_NAME:
                raise ValueError(f"raw array dtype {dtype!r} is not supported")
            shape = tuple(int(d) for d in spec.get("shape") or ())
            n, code = _count(shape), _BY_NAME[dtype]
            if len(raw) != n * struct.calcsize("<" + code):
                raise ValueError(f"raw array holds {len(raw)} bytes, shape {list(shape)} {dtype} needs "
                                 f"{n * struct.calcsize('<' + code)}")
            arr = Array(shape, dtype, list(struct.unpack(f"<{n}{code}", raw)))
        else:
            with zipfile.ZipFile(io.BytesIO(raw)) as z:
                names = z.namelist()
                member = f"{spec.get('key')}.npy"
                if member not in names:
                    if len(names) != 1:
                        raise ValueError(f"npz has no member {member!r} (members {names})")
                    member = names[0]
                arr = _npy(z.read(member))
    if "shape" in spec and list(arr.shape) != [int(d) for d in spec["shape"]]:
        raise ValueError(f"decoded shape {list(arr.shape)} differs from the declared {list(spec['shape'])}")
    if len(arr.values) != _count(arr.shape):
        raise ValueError(f"{len(arr.values)} elements do not fill shape {list(arr.shape)}")
    return arr


def _is_encoded_array(value: Any) -> bool:
    return isinstance(value, Mapping) and value.get("format") in ARRAY_FORMATS and "data" in value


# ---------------------------------------------------------------------------------------------- comparison

def find_reference(sample: PathLike, *keys: Optional[str]) -> Optional[Path]:
    """The stored CPU-reference body of ``sample``: the first existing ``<stem>.<key>.reference.json`` for the
    given keys (e.g. the serve profile, then the model variant), else ``<stem>.reference.json``; None if absent."""
    p = Path(sample)
    names = [f"{p.stem}.{k}.reference.json" for k in keys if k and _SAFE_KEY.fullmatch(str(k))]
    for name in names + [f"{p.stem}.reference.json"]:
        if (p.parent / name).is_file():
            return p.parent / name
    return None


def _label(row: Mapping[str, Any]) -> Any:
    return row.get("label_id", row.get("label"))


def _iou(a: Sequence[float], b: Sequence[float]) -> float:
    ix = max(0.0, min(a[2], b[2]) - max(a[0], b[0]))
    iy = max(0.0, min(a[3], b[3]) - max(a[1], b[1]))
    inter = ix * iy
    union = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
    return inter / union if union > 0 else 0.0


def _greedy_match(test: List[Mapping[str, Any]], ref: List[Mapping[str, Any]],
                  cost: Callable[[Mapping[str, Any], Mapping[str, Any]], float],
                  accept: Callable[[float], bool]) -> List[Tuple[int, int, float]]:
    """Reference rows in descending score order (stable) each take the cheapest unused same-label test row;
    returns ``(test index, reference index, cost)`` of the accepted pairs."""
    used = [False] * len(test)
    pairs = []
    for i in sorted(range(len(ref)), key=lambda k: -float(ref[k].get("score", 0.0))):
        best, best_cost = -1, math.inf
        for j, row in enumerate(test):
            if not used[j] and _label(row) == _label(ref[i]):
                c = cost(row, ref[i])
                if c < best_cost:
                    best, best_cost = j, c
        if best >= 0 and accept(best_cost):
            used[best] = True
            pairs.append((best, i, best_cost))
    return pairs


def _compare_detections(test: Any, ref: List[Mapping[str, Any]],
                        g: Mapping[str, Any]) -> Tuple[Dict[str, Any], List[str]]:
    if not isinstance(test, list):
        return {}, ["the response has no 'detections' list"]
    rows = ref + test
    if any("center" in d for d in rows):
        kind = "3d"
        pairs = _greedy_match(test, ref, lambda t, r: math.hypot(t["center"][0] - r["center"][0],
                                                                 t["center"][1] - r["center"][1]),
                              lambda c: c <= g["max_center_dist"])
    elif any("box_xyxy" in d for d in rows):
        kind = "2d"
        pairs = _greedy_match(test, ref, lambda t, r: 1.0 - _iou(t["box_xyxy"], r["box_xyxy"]),
                              lambda c: 1.0 - c >= g["min_iou"])
    elif not rows:  # nothing detected on either side (e.g. a synthetic sample): agreement
        kind, pairs = "none", []
    else:
        return {}, ["detections carry neither 'center' nor 'box_xyxy'"]
    score_err = max((abs(float(test[j]["score"]) - float(ref[i]["score"])) for j, i, _ in pairs), default=0.0)
    m: Dict[str, Any] = {"kind": kind, "n": len(test), "n_ref": len(ref), "matched": len(pairs),
                         "recall": len(pairs) / len(ref) if ref else 1.0,
                         "precision": len(pairs) / len(test) if test else 1.0, "score_err_max": score_err}
    if kind == "3d":
        m["center_err_max"] = max((c for _, _, c in pairs), default=0.0)
    elif kind == "2d":
        m["iou_min"] = min((1.0 - c for _, _, c in pairs), default=1.0)
    fails = []
    if m["recall"] < g["min_recall"]:
        fails.append(f"detections recall {m['recall']:.3f} < {g['min_recall']} ({len(pairs)}/{len(ref)} reference "
                     "detections matched)")
    if m["precision"] < g["min_precision"]:
        fails.append(f"detections precision {m['precision']:.3f} < {g['min_precision']} ({len(pairs)}/{len(test)} "
                     "served detections matched)")
    if score_err > g["max_score_err"]:
        fails.append(f"detections max |score difference| {score_err:.4f} > {g['max_score_err']}")
    return m, fails


def _compare_trajectory(body: Mapping[str, Any], ref: Mapping[str, Any],
                        g: Mapping[str, Any]) -> Tuple[Dict[str, Any], List[str]]:
    test, want = body.get("trajectory"), ref["trajectory"]
    if not isinstance(test, list):
        return {}, ["the response has no 'trajectory'"]
    if len(test) != len(want):
        return {"n": len(test), "n_ref": len(want)}, [f"trajectory has {len(test)} poses, the reference {len(want)}"]
    columns = list(ref.get("columns") or ["x", "y"])
    if body.get("columns", columns) != columns:
        return {}, [f"trajectory columns {body.get('columns')} differ from the reference {columns}"]
    ix, iy = (columns.index("x"), columns.index("y")) if {"x", "y"} <= set(columns) else (0, 1)
    disp = [math.hypot(a[ix] - b[ix], a[iy] - b[iy]) for a, b in zip(test, want)]
    m = {"n": len(test), "ade": sum(disp) / len(disp) if disp else 0.0, "fde": disp[-1] if disp else 0.0}
    fails = []
    if not m["ade"] <= g["max_ade"]:
        fails.append(f"trajectory ADE {m['ade']:.4f} > {g['max_ade']}")
    if not m["fde"] <= g["max_fde"]:
        fails.append(f"trajectory FDE {m['fde']:.4f} > {g['max_fde']}")
    return m, fails


def _compare_array(name: str, test: Any, ref: Mapping[str, Any],
                   g: Mapping[str, Any]) -> Tuple[Dict[str, Any], List[str]]:
    if not _is_encoded_array(test):
        return {}, [f"the response has no encoded array {name!r}"]
    a, b = decode_array(test), decode_array(ref)
    if a.shape != b.shape:
        return {"shape": list(a.shape), "shape_ref": list(b.shape)}, [
            f"{name}: shape {list(a.shape)} differs from the reference {list(b.shape)}"]
    if a.is_float or b.is_float:
        err = 0.0
        for x, y in zip(a.values, b.values):
            if x == y:  # equal values, equal infinities included
                continue
            d = abs(float(x) - float(y))
            if d != d:  # NaN on either side
                err = math.nan
                break
            err = max(err, d)
        limit = g["max_abs_err"]
        fails = [f"{name}: max |difference| {err:.6g} > {limit}"] if limit is not None and not err <= limit else []
        return {"max_abs_err": err, "n": len(b.values)}, fails
    equal = sum(1 for x, y in zip(a.values, b.values) if x == y)
    agreement = equal / len(b.values) if b.values else 1.0
    fails = []
    if agreement < g["min_label_agreement"]:
        fails.append(f"{name}: agreement {agreement:.4f} < {g['min_label_agreement']} "
                     f"({len(b.values) - equal} of {len(b.values)} elements differ)")
    return {"agreement": agreement, "n": len(b.values)}, fails


def compare_with_reference(body: Mapping[str, Any], reference: Mapping[str, Any],
                           gates: Optional[Mapping[str, Optional[float]]] = None
                           ) -> Tuple[Dict[str, Dict[str, Any]], List[str]]:
    """Compare a served ``/predict`` body with the stored CPU-reference body of the same input (module docstring).
    ``gates`` override :data:`DEFAULT_GATES` (unknown names raise ``ValueError``). Returns ``(metrics per compared
    key, failures)``; a body that cannot be decoded is a failure, not an exception."""
    unknown = sorted(set(gates or {}) - set(DEFAULT_GATES))
    if unknown:
        raise ValueError(f"unknown reference gates {unknown}; known: {sorted(DEFAULT_GATES)}")
    g = {**DEFAULT_GATES, **dict(gates or {})}
    metrics: Dict[str, Dict[str, Any]] = {}
    fails = [f"{k} is {body.get(k)!r}, the reference has {reference[k]!r}" for k in ("model", "frame_id")
             if k in reference and body.get(k) != reference[k]]
    checks: List[Tuple[str, Callable[[], Tuple[Dict[str, Any], List[str]]]]] = []
    if "detections" in reference:
        checks.append(("detections", lambda: _compare_detections(body.get("detections"), reference["detections"], g)))
    if "trajectory" in reference:
        checks.append(("trajectory", lambda: _compare_trajectory(body, reference, g)))
    for key, value in reference.items():
        if key != "arrays" and _is_encoded_array(value):
            checks.append((key, lambda k=key, v=value: _compare_array(k, body.get(k), v, g)))
    if not checks:
        fails.append("the reference holds nothing to compare (no detections, trajectory or encoded arrays)")
    for key, check in checks:
        try:
            metrics[key], more = check()
        except (ValueError, KeyError, TypeError, IndexError, struct.error, zlib.error, zipfile.BadZipFile) as e:
            metrics[key], more = {}, [f"{key}: cannot compare ({type(e).__name__}: {e})"]
        fails += more
    return metrics, fails


def describe_metrics(metrics: Mapping[str, Mapping[str, Any]]) -> str:
    """One compact line for the smoke verdict, e.g. ``detections 5/5 recall=1.000 precision=1.000 dscore=0.0040``."""
    parts = []
    for key, m in metrics.items():
        if "recall" in m:
            parts.append(f"{key} {m['matched']}/{m['n_ref']} recall={m['recall']:.3f} "
                         f"precision={m['precision']:.3f} dscore={m['score_err_max']:.4f}")
        elif "ade" in m:
            parts.append(f"{key} ade={m['ade']:.4f} fde={m['fde']:.4f}")
        elif "agreement" in m:
            parts.append(f"{key} agree={m['agreement']:.4f}")
        elif "max_abs_err" in m:
            parts.append(f"{key} max_abs={m['max_abs_err']:.4g}")
        else:
            parts.append(f"{key} -")
    return "; ".join(parts) or "-"


# ----------------------------------------------------------------------------------------- package pins

def pinned_config(manifest: PathLike, profile: Optional[str] = None) -> Dict[str, Any]:
    """What a staged package (``<out>/<name>/tt_kernel_manifest.json``) pins for one serve profile (default: the
    package's default profile): ``{"profile": name, "env": serve.env with the profile's env on top (tt-model's
    merge rule), "weights": {"repo", "revision"}}``."""
    m = load_json(manifest)
    c = m.get("container") or {}
    profiles = c.get("serve_profiles") or [{"name": "default", "env": {}}]
    name = profile or c.get("default_profile") or profiles[0].get("name")
    chosen = [p for p in profiles if p.get("name") == name]
    if not chosen:
        raise ValueError(f"no serve profile {name!r} in {manifest}; available: {[p.get('name') for p in profiles]}")
    env = dict((c.get("serve") or {}).get("env") or {})
    env.update({k: v for k, v in (chosen[0].get("env") or {}).items() if v is not None})
    weights = m.get("weights") or {}
    return {"profile": name, "env": {k: str(v) for k, v in env.items()},
            "weights": {"repo": weights.get("repo_id"), "revision": weights.get("revision")}}


def check_pinned(info: Mapping[str, Any], pinned: Mapping[str, Any], prefix: str) -> List[str]:
    """``/info`` runs what the package pins: ``<prefix>_DISPATCH`` (eth | worker), ``<prefix>_NUM_CQS``,
    ``<prefix>_VARIANT`` and the weights revision. Pins that are absent are not checked. Returns failures."""
    env = pinned.get("env") or {}
    device = info.get("device") or {}
    profile = pinned.get("profile")
    fails = []
    dispatch = (env.get(f"{prefix}_DISPATCH") or "").strip().lower()
    if dispatch in ("eth", "worker") and device.get("dispatch") != dispatch:
        fails.append(f"dispatch is {device.get('dispatch')!r}, profile {profile!r} pins {prefix}_DISPATCH={dispatch}")
    cqs = (env.get(f"{prefix}_NUM_CQS") or "").strip()
    if cqs and str(device.get("num_command_queues")) != cqs:
        fails.append(f"num_command_queues is {device.get('num_command_queues')!r}, profile {profile!r} pins "
                     f"{prefix}_NUM_CQS={cqs}")
    variant = (env.get(f"{prefix}_VARIANT") or "").strip()
    if variant and info.get("variant") != variant:
        fails.append(f"variant is {info.get('variant')!r}, profile {profile!r} pins {prefix}_VARIANT={variant}")
    revision = (pinned.get("weights") or {}).get("revision")
    served = (info.get("weights") or {}).get("revision")
    if revision and served != revision:
        fails.append(f"weights revision is {served!r}, the package pins {revision!r}")
    return fails