"""
Reference live client for the Dorayaki recognition server.

Streams frames over WebSocket /v1/stream and prints what the server sees. This
is the documented example for capture sources that are not a browser: a cart
camera, a wearable camera, a rail camera, a robot.

Usage (from recognization/):
    .venv/bin/python code/stream_client.py --source image.jpg --pan 12 --seconds 20
    .venv/bin/python code/stream_client.py --source 0                  # webcam index
    .venv/bin/python code/stream_client.py --source clip.mp4 --seconds 30
    .venv/bin/python code/stream_client.py --url ws://host:8080/v1/stream --token SECRET

Frame protocol: 4-byte big-endian frame id + JPEG bytes.
"""

from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
import time
from typing import Any, Dict, List, Optional

current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
if os.path.exists(os.path.join(project_root, "stock")):
    os.chdir(project_root)
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

import cv2  # noqa: E402

from websockets.sync.client import connect  # noqa: E402


def percentile(values: List[float], pct: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    index = min(len(ordered) - 1, max(0, int(round((pct / 100.0) * (len(ordered) - 1)))))
    return ordered[index]


class FrameSource:
    """Yields BGR frames from a camera, a video file, or a still image pan."""

    def __init__(self, source: str, *, width: int, pan_step: int, pan_width: int = 1280):
        self.source = source
        self.width = width
        self.pan_step = pan_step
        self.pan_width = pan_width
        self.capture: Optional[cv2.VideoCapture] = None
        self.image = None
        self.pan_x = 0
        self.pan_direction = 1
        # Where the last frame was cut from the source image. (0, 0) for a
        # camera or a video, whose frames are the whole image.
        self.offset = (0, 0)
        # The source-image rects the window showed, so an evaluator can restrict
        # its denominator to the facings this run could actually have seen.
        # A list, because a panned run shows a different rect each frame and the
        # reachable set is their union. Both axes matter: a 1280-wide window
        # over a 1500-wide photo is x-limited as well as y-limited.
        self.windows: List[Tuple[int, int, int, int]] = []
        if source.isdigit():
            self.capture = cv2.VideoCapture(int(source))
        elif os.path.isfile(source) and source.lower().endswith((".mp4", ".mov", ".avi", ".mkv", ".webm")):
            self.capture = cv2.VideoCapture(source)
        else:
            self.image = cv2.imread(source)
            if self.image is None:
                raise SystemExit(f"cannot read source {source}")

    def frame(self, index: int):
        """One frame, plus where it was cut from the source image.

        The offset is what lets an evaluator map a track's box back to the
        source image's coordinates: a label measured on a pan window is only
        comparable to ground truth once the window's origin is added back.
        """
        if self.image is not None:
            height, width = self.image.shape[:2]
            window_w = min(self.pan_width, width)
            window_h = min(int(window_w * 9 / 16), height)
            max_x = max(0, width - window_w)
            x = min(self.pan_x, max_x)
            y = max(0, (height - window_h) // 2)
            window = self.image[y:y + window_h, x:x + window_w]
            self.offset = (x, y)
            self.windows.append((x, y, x + window_w, y + window_h))
            self.pan_x += self.pan_direction * self.pan_step
            if self.pan_x >= max_x or self.pan_x <= 0:
                self.pan_direction *= -1
            return window
        if self.capture is not None:
            ok, frame = self.capture.read()
            if not ok:  # loop video files
                self.capture.set(cv2.CAP_PROP_POS_FRAMES, 0)
                ok, frame = self.capture.read()
            if not ok:
                return None
            return frame
        return None

    def release(self) -> None:
        if self.capture is not None:
            self.capture.release()


def encode(frame, width: int, quality: int) -> bytes:
    height, current_width = frame.shape[:2]
    if width and current_width > width:
        scale = width / float(current_width)
        frame = cv2.resize(frame, (int(width), int(round(height * scale))))
    ok, buffer = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, quality])
    if not ok:
        raise RuntimeError("JPEG encode failed")
    return buffer.tobytes()


def main() -> None:
    parser = argparse.ArgumentParser(description="Dorayaki live stream client")
    parser.add_argument("--url", default="ws://127.0.0.1:8080/v1/stream")
    parser.add_argument("--token", default=os.environ.get("DORAYAKI_API_TOKEN"))
    parser.add_argument("--source", default="image.jpg", help="camera index, video path, or image path")
    parser.add_argument("--pan", type=int, default=12, help="pixels per frame when panning a still image")
    parser.add_argument("--seconds", type=float, default=20.0)
    parser.add_argument("--json", action="store_true", help="print a JSON summary instead of the table")
    parser.add_argument("--tracks", default=None,
                        help="write per-track labels and their boxes (source-image "
                             "coordinates) as JSON, for scoring against ground truth")
    parser.add_argument("--quiet", action="store_true", help="only print the final summary")
    parser.add_argument("--width", type=int, default=1280, help="send width")
    parser.add_argument("--quality", type=int, default=75)
    parser.add_argument("--max-inflight", type=int, default=1)
    parser.add_argument("--pace", choices=["send", "result"], default="send",
                        help="'send': keep the socket busy (frames drop server-side); "
                             "'result': send the next frame only after a result, so the "
                             "pan stays continuous (the mode to use for stability checks)")
    parser.add_argument("--method", default=None, help="surf | orb | brisk | hybrid")
    args = parser.parse_args()

    url = args.url
    headers = {"Authorization": f"Bearer {args.token}"} if args.token else None

    source = FrameSource(args.source, width=args.width, pan_step=args.pan)
    rtts: List[float] = []
    detect_times: List[float] = []
    frames_sent = 0
    results = 0
    drops = 0
    first_label: Dict[int, float] = {}
    label_latencies: List[float] = []
    labels: Dict[int, str] = {}
    pending_frame: Optional[int] = None
    # Per track: the label it settled on, the box that label was measured on,
    # and the pan offset of the frame that box lives in. Written out with
    # --tracks so an evaluator can match labels to ground truth.
    track_records: Dict[int, Dict[str, Any]] = {}
    frame_offsets: Dict[int, Tuple[int, int]] = {}
    label_changes = 0
    unexplained_changes = 0
    relabel_seen: Dict[int, int] = {}
    restart_seen: Dict[int, int] = {}
    inherited_seen: Dict[int, bool] = {}
    seen_tracks: set = set()
    last_summary: Dict[str, Any] = {}
    started = time.perf_counter()
    in_flight: Dict[int, float] = {}
    pending_frame: Optional[int] = None

    with connect(url, max_size=8 * 1024 * 1024, additional_headers=headers) as socket:
        hello = json.loads(socket.recv(timeout=10))
        if hello.get("type") != "hello":
            raise SystemExit(f"unexpected first message: {hello}")
        if not args.quiet:
            print(f"hello: session {hello['session_id']} model {hello['model']['version']} "
                  f"device {hello['model']['device']} config {hello['config']}")
        if args.method:
            socket.send(json.dumps({"type": "config", "method": args.method}))
            ack = json.loads(socket.recv(timeout=10))
            if not args.quiet:
                print(f"config ack: {ack.get('config')}")

        def consume(message: Dict[str, Any]) -> None:
            """Fold one result into the running metrics."""
            nonlocal results, last_summary, label_changes, unexplained_changes
            if message.get("type") != "result":
                if not args.quiet:
                    print(f"server message: {message}")
                return
            received_at = time.perf_counter()
            sent_at = in_flight.pop(message["frame_id"], None)
            if sent_at is not None:
                rtts.append((received_at - sent_at) * 1000.0)
            detect_times.append(message["timing_ms"]["detect"])
            results += 1
            last_summary = message["summary"]
            now = time.perf_counter()
            for track in message["tracks"]:
                track_id = track["track_id"]
                seen_tracks.add(track_id)
                identity = track["identity"]
                if identity["state"] != "recognized":
                    continue
                sku = identity["sku_id"]
                relabels = track.get("relabels", 0)
                restarts = track.get("restarts", 0)
                if restarts > restart_seen.get(track_id, 0):
                    # The id was restarted: earlier labels describe another object.
                    labels.pop(track_id, None)
                    first_label.pop(track_id, None)
                inherited = bool(track.get("inherited"))
                if track_id not in first_label:
                    first_label[track_id] = now
                    if not inherited:
                        label_latencies.append((now - started) * 1000.0)
                # Only a measured-to-measured change on the same object is a
                # label change; inherited labels describe the same product
                # under a new track id.
                if track_id in labels and labels[track_id] != sku:
                    if not inherited and not inherited_seen.get(track_id):
                        label_changes += 1
                        if relabels <= relabel_seen.get(track_id, 0):
                            unexplained_changes += 1
                labels[track_id] = sku
                # The box the label is about, in source-image coordinates. The
                # offset comes from the frame the box was measured on, not the
                # frame this result arrived with: the pan has moved on since.
                offset = frame_offsets.get(track.get("label_frame_id"), (0, 0))
                label_box = track.get("label_box") or track.get("box") or []
                if len(label_box) == 4:
                    track_records[track_id] = {
                        "track_id": track_id,
                        "sku_id": sku,
                        "label_box": [
                            label_box[0] + offset[0], label_box[1] + offset[1],
                            label_box[2] + offset[0], label_box[3] + offset[1],
                        ],
                        "label_frame_id": track.get("label_frame_id"),
                        "result_frame_id": message["frame_id"],
                        "inherited": inherited,
                        "restarts": restarts,
                    }
                inherited_seen[track_id] = inherited
                relabel_seen[track_id] = relabels
                restart_seen[track_id] = restarts
            if not args.quiet and results % 10 == 0:
                print(f"f{message['seq']:04d} tracks={len(message['tracks']):3d} "
                      f"recognized={last_summary['recognized']:3d} "
                      f"unconfirmed={last_summary.get('unconfirmed', 0):3d} "
                      f"detect={message['timing_ms']['detect']:5.1f}ms "
                      f"rtt={rtts[-1]:6.1f}ms fps={results / (now - started):5.2f}")

        frame_id = 0
        index = 0
        while time.perf_counter() - started < args.seconds:
            frame = source.frame(index)
            if frame is None:
                break
            index += 1
            frame_id = (frame_id + 1) & 0xFFFFFFFF
            frame_offsets[frame_id] = source.offset
            payload = encode(frame, args.width, args.quality)
            if len(in_flight) >= args.max_inflight:
                if pending_frame is not None:
                    drops += 1
                pending_frame = frame_id
                socket.send(frame_id.to_bytes(4, "big") + payload, text=False)
                frames_sent += 1
            else:
                socket.send(frame_id.to_bytes(4, "big") + payload, text=False)
                frames_sent += 1
            in_flight[frame_id] = time.perf_counter()
            if args.pace == "result":
                # One frame in flight: the next frame is captured only after this
                # one was answered, so the pan stays continuous.
                while in_flight:
                    try:
                        consume(json.loads(socket.recv(timeout=5)))
                    except TimeoutError:
                        print("timeout waiting for a result")
                        break
                continue

            try:
                consume(json.loads(socket.recv(timeout=5)))
            except TimeoutError:
                print("timeout waiting for a result")
                break

    source.release()
    elapsed = time.perf_counter() - started
    summary = {
        "seconds": round(elapsed, 1),
        "frames_sent": frames_sent,
        "results": results,
        "results_per_s": round(results / elapsed, 2) if elapsed else 0,
        "frames_dropped_by_client": drops,
        "rtt_ms_p50": round(statistics.median(rtts), 1) if rtts else None,
        "rtt_ms_p95": round(percentile(rtts, 95), 1) if rtts else None,
        "server_detect_ms_p50": round(statistics.median(detect_times), 1) if detect_times else None,
        "server_detect_ms_p95": round(percentile(detect_times, 95), 1) if detect_times else None,
        "tracks_seen": len(seen_tracks),
        "tracks_labelled": len(labels),
        "label_changes": label_changes,
        "unexplained_label_changes": unexplained_changes,
        "first_label_ms_p50": round(statistics.median(label_latencies), 0) if label_latencies else None,
        "final_summary": last_summary,
    }
    if args.tracks:
        with open(args.tracks, "w") as handle:
            json.dump({
                "windows": [list(w) for w in source.windows] or None,
                "source": args.source,
                "pace": args.pace,
                "seconds": round(elapsed, 1),
                "tracks": sorted(track_records.values(), key=lambda row: row["track_id"]),
            }, handle, indent=2)
    if args.json:
        print(json.dumps(summary, indent=2))
    else:
        print("\n--- stream summary ---")
        for key, value in summary.items():
            if key != "final_summary":
                print(f"{key:24s} {value}")
        print(f"{'final_summary':24s} {json.dumps(last_summary, ensure_ascii=False)}")
    if results == 0:
        sys.exit(1)


if __name__ == "__main__":
    main()