#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""C08 client: build / send ``POST /predict`` requests and run smoke checks. Standard library only.

This module imports nothing from ttaw (no relative imports), so it works three ways: as ``<pkg>.ttaw.server.client``,
as a script, and loaded by path from a bundle's ``smoke_test.py`` running under the host's system ``python3``
(``container_smoke.sh``)::

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

CLI::

    python client.py --points sample.npz --out req.json                    # write the request
    python client.py --points sample.npz --url http://127.0.0.1:20000      # send it, print the response
    python client.py --image CAM_FRONT=f.jpg --calib-preset rig --url ...  # cameras
    python client.py --inputs scene.npz --url ...                          # planner
"""
from __future__ import annotations

import argparse
import base64
import json
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

__all__ = ["POINT_SUFFIXES", "b64file", "build_request", "post", "get_json", "wait_ready", "check_service",
           "check_bad_request", "smoke_line", "main"]

POINT_SUFFIXES = {".bin": "bin", ".npy": "npy", ".npz": "npz", ".pcd": "pcd"}


def b64file(path: Any) -> str:
    return base64.b64encode(Path(path).read_bytes()).decode("ascii")


def build_request(points: Optional[str] = None, fields: Optional[Sequence[str]] = None, images: Iterable[str] = (),
                  calib: Optional[str] = None, calib_preset: Optional[str] = None, inputs: Optional[str] = None,
                  params: Optional[Dict[str, Any]] = None, stream_id: Optional[str] = None, reset: bool = False,
                  output_format: str = "json") -> Dict[str, Any]:
    """A ``/predict`` body from files: ``points`` path (.bin needs ``fields``), ``images`` as ``NAME=PATH``,
    ``calib`` JSON file or ``calib_preset`` name, ``inputs`` .npz (planner), ``params``, stream id / reset."""
    req: Dict[str, Any] = {}
    if points:
        p = Path(points)
        fmt = POINT_SUFFIXES.get(p.suffix.lower())
        if fmt is None:
            raise ValueError(f"points: unknown suffix {p.suffix!r} (use .bin / .npy / .npz / .pcd)")
        req["points"] = {"format": fmt, "data": b64file(p)}
        if fields:
            req["points"]["fields"] = list(fields)
    images = list(images)
    if images:
        req["images"] = []
        for spec in images:
            name, _, path = spec.partition("=")
            if not path:
                raise ValueError(f"image expects NAME=PATH, got {spec!r}")
            req["images"].append({"camera": name, "data": b64file(path)})
    if calib:
        req["calibration"] = json.loads(Path(calib).read_text())
    elif calib_preset:
        req["calibration"] = {"preset": calib_preset}
    if inputs:
        req["inputs"] = {"format": "npz", "data": b64file(inputs)}
    if params:
        req["params"] = dict(params)
    if stream_id is not None or reset:
        req["stream"] = {"id": stream_id or "default", "reset": bool(reset)}
    if output_format != "json":
        req["output_format"] = output_format
    return req


def post(url: str, payload: Dict[str, Any], timeout: float = 120.0) -> Tuple[int, Dict[str, Any]]:
    """POST ``payload`` to ``<url>/predict`` -> ``(status, json body)`` (HTTP errors are returned, not raised)."""
    data = json.dumps(payload).encode("utf-8")
    r = urllib.request.Request(url.rstrip("/") + "/predict", data=data, headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(r, timeout=timeout) as resp:
            return resp.status, json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", "replace")
        try:
            return e.code, json.loads(body)
        except ValueError:
            return e.code, {"detail": body}


def get_json(url: str, timeout: float = 10.0) -> Dict[str, Any]:
    with urllib.request.urlopen(url, timeout=timeout) as r:
        return json.loads(r.read().decode("utf-8"))


def wait_ready(base: str, wait_s: float = 0.0, poll_s: float = 5.0) -> Dict[str, Any]:
    """Poll ``/health`` until ``ok`` / ``error`` or ``wait_s`` elapses; returns the last health body."""
    deadline = time.time() + wait_s
    last: Dict[str, Any] = {}
    while True:
        try:
            last = get_json(base.rstrip("/") + "/health")
            if last.get("status") in ("ok", "error"):
                return last
        except (urllib.error.URLError, OSError, ValueError) as e:
            last = {"status": f"unreachable ({e})"}
        if time.time() >= deadline:
            return last
        time.sleep(poll_s)


def check_service(base: str, *, expect_dispatch: Optional[str] = "eth",
                  expect_grid: Optional[str] = "12x10") -> Tuple[Dict[str, Any], List[str]]:
    """``/info`` and ``/v1/models`` answer; ``/info`` reports the pinned dispatch and grid (PLAN.md 0.3 item 6:
    the smoke must assert ETH dispatch and 12x10, not just print them). Returns ``(info, failures)``."""
    fails: List[str] = []
    try:
        info = get_json(base.rstrip("/") + "/info")
        models = get_json(base.rstrip("/") + "/v1/models")
    except (urllib.error.URLError, OSError, ValueError) as e:
        return {}, [f"/info or /v1/models: {e}"]
    if not models.get("data"):
        fails.append("/v1/models is empty")
    device = info.get("device") or {}
    if expect_dispatch and device.get("dispatch") != expect_dispatch:
        fails.append(f"dispatch is {device.get('dispatch')!r}, expected {expect_dispatch!r}")
    if expect_grid and device.get("grid") != expect_grid:
        fails.append(f"grid is {device.get('grid')!r}, expected {expect_grid!r}")
    return info, fails


def check_bad_request(base: str, payload: Optional[Dict[str, Any]] = None) -> Optional[str]:
    """A malformed request must answer 400 (not 500). Returns a failure message or None."""
    code, _ = post(base, payload or {"points": {"format": "npy", "data": "not base64!"}})
    return None if code == 400 else f"malformed request answered {code}, expected 400"


def smoke_line(model: str, summary: str, failures: Sequence[str]) -> str:
    """The one-line verdict container_smoke.sh greps: ``PASS <model>: ...`` / ``FAIL <model>: ... :: reasons``."""
    if failures:
        return f"FAIL {model}: {summary} :: " + "; ".join(list(failures)[:5])
    return f"PASS {model}: {summary}"


def _param(text: str) -> Tuple[str, Any]:
    k, _, v = text.partition("=")
    try:
        return k, json.loads(v)
    except ValueError:
        return k, v


def main(argv: Optional[Sequence[str]] = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--points")
    ap.add_argument("--fields", help="comma list, e.g. x,y,z,intensity (needed for .bin)")
    ap.add_argument("--image", action="append", default=[], help="NAME=PATH, repeat per camera")
    ap.add_argument("--calib", help="calibration JSON file")
    ap.add_argument("--calib-preset", help="a preset name listed by /info")
    ap.add_argument("--inputs", help=".npz of named input tensors (planner)")
    ap.add_argument("--param", action="append", default=[], help="k=v (JSON value), repeat")
    ap.add_argument("--stream-id")
    ap.add_argument("--reset", action="store_true")
    ap.add_argument("--output-format", default="json", choices=["json", "npz"])
    ap.add_argument("--out", help="write the request JSON here")
    ap.add_argument("--url", help="POST it to this server and print the response")
    a = ap.parse_args(argv)
    req = build_request(a.points, a.fields.split(",") if a.fields else None, a.image, a.calib, a.calib_preset,
                        a.inputs, dict(_param(p) for p in a.param), a.stream_id, a.reset, a.output_format)
    if a.out:
        Path(a.out).write_text(json.dumps(req))
    if a.url:
        code, body = post(a.url, req)
        print(json.dumps(body, indent=1)[:20000])
        return 0 if code == 200 else 1
    if not a.out:
        json.dump(req, sys.stdout)
    return 0


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