"""
Synthetic camera for demos: a store employee walking a shelf.

A simulator has no camera, and a still image makes a poor demo - nothing moves, so
nothing new is ever recognised. This module turns the demo shelf photos into a
virtual walk: a window pans left to right along a row, steps up to the next row,
and pans back the other way (serpentine, ascending), exactly the way a person
sweeps a shelf with a tablet.

The client drives only `step`; all geometry lives here, so the web client, the
phone app, and the reference client see the same walk. `step` wraps, so an
unattended demo loops forever.

Demo-only: this is a capture source, not a perception feature. `/v1/demo/*` is
documented as such.
"""

from __future__ import annotations

import math
import os
import threading
import time
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple

DEFAULT_WINDOW = (960, 540)
DEFAULT_QUALITY = 65
# How far the window advances per step, as a fraction of its own size. This is
# the demo's most important number: advance too fast and a product is in view for
# one or two frames, so the classifier never gets a chance and nothing is ever
# labelled (measured: 0.22 gave 9 labelled tracks in 20 s; the shelf swept past
# faster than a person could read it). 0.06 keeps a facing in view for roughly a
# second at demo frame rates - a deliberate, human sweep.
PAN_OVERLAP = 0.06
ROW_OVERLAP = 0.25
# Vertical steps spent moving between rows (the "look up, then keep going" beat).
TRANSITION_STEPS = 6


class DemoSourceError(ValueError):
    """Unknown sample or an impossible window."""


@dataclass
class Bay:
    """One shelf photo in the walk, with its geometry."""

    sample_id: str
    path: str
    width: int
    height: int
    steps: int


class DemoCamera:
    """Plans and renders the virtual walk over one or more shelf photos."""

    def __init__(self, samples: Dict[str, Dict[str, str]], window: Tuple[int, int] = DEFAULT_WINDOW,
                 pan_overlap: float = PAN_OVERLAP):
        self._samples = samples
        self._window = window
        # How far the window advances per step. The demo's pacing knob: a smaller
        # value is a slower, more deliberate sweep (more steps, a facing in view
        # for longer). Clamped so a client cannot ask for a frozen or teleporting
        # camera.
        self.pan_overlap = max(0.005, min(0.5, float(pan_overlap)))
        self._lock = threading.Lock()
        self._images: Dict[str, Any] = {}
        self._sizes: Dict[str, Tuple[int, int]] = {}
        self._plans: Dict[Tuple[str, ...], List[Bay]] = {}

    # ------------------------------------------------------------------ planning
    def _size(self, sample_id: str) -> Tuple[int, int]:
        cached = self._sizes.get(sample_id)
        if cached is not None:
            return cached
        item = self._samples.get(sample_id)
        if item is None or not os.path.exists(item["path"]):
            raise DemoSourceError(f"unknown sample '{sample_id}'")
        import cv2

        image = cv2.imread(item["path"])
        if image is None:
            raise DemoSourceError(f"cannot read sample '{sample_id}'")
        with self._lock:
            self._images[sample_id] = image
            self._sizes[sample_id] = (int(image.shape[1]), int(image.shape[0]))
        return self._sizes[sample_id]

    def _image(self, sample_id: str) -> Any:
        if sample_id not in self._images:
            self._size(sample_id)
        return self._images[sample_id]

    def steps_for(self, sample_id: str, window: Tuple[int, int]) -> int:
        width, height = self._size(sample_id)
        win_w, win_h = window
        bands = self._bands(height, win_h)
        per_band = self._pan_steps(width, win_w)
        # Every band pans, and every band after the first needs the step-up.
        return bands * per_band + max(0, bands - 1) * TRANSITION_STEPS

    def _bands(self, image_height: int, window_height: int) -> int:
        if window_height >= image_height:
            return 1
        stride = max(1, int(window_height * (1.0 - ROW_OVERLAP)))
        return max(1, int(math.ceil((image_height - window_height) / stride)) + 1)

    def _pan_steps(self, image_width: int, window_width: int) -> int:
        if window_width >= image_width:
            return 1
        stride = max(1, int(window_width * self.pan_overlap))
        return max(1, int(math.ceil((image_width - window_width) / stride)) + 1)

    def plan(self, sample_ids: Sequence[str], window: Optional[Tuple[int, int]] = None) -> Dict[str, Any]:
        window = window or self._window
        key = (tuple(sample_ids), window)
        bays = self._plans.get(key)
        if bays is None:
            bays = []
            for sample_id in sample_ids:
                width, height = self._size(sample_id)
                bays.append(
                    Bay(
                        sample_id=sample_id,
                        path=self._samples[sample_id]["path"],
                        width=width,
                        height=height,
                        steps=self.steps_for(sample_id, window),
                    )
                )
            self._plans[key] = bays
        total = sum(bay.steps for bay in bays)
        return {
            "samples": [bay.sample_id for bay in bays],
            "window": {"width": window[0], "height": window[1]},
            "bays": [
                {"sample": bay.sample_id, "width": bay.width, "height": bay.height,
                 "steps": bay.steps, "first_step": sum(b.steps for b in bays[:index])}
                for index, bay in enumerate(bays)
            ],
            "steps": total,
            "row_overlap": ROW_OVERLAP,
            "pan_overlap": self.pan_overlap,
        }

    # ------------------------------------------------------------------ rendering
    def window_at(
        self, sample_ids: Sequence[str], step: int, window: Optional[Tuple[int, int]] = None
    ) -> Dict[str, Any]:
        """Where the virtual camera is at this step (wraps around the walk)."""
        window = window or self._window
        plan = self.plan(sample_ids, window)
        total = plan["steps"]
        step = int(step) % max(1, total)
        offset = 0
        for index, bay in enumerate(plan["bays"]):
            if step < offset + bay["steps"]:
                local = step - offset
                x, y = self._position(local, bay["width"], bay["height"], window)
                return {
                    "sample": bay["sample"],
                    "bay_index": index,
                    "step": step,
                    "steps": total,
                    "local_step": local,
                    "bay_steps": bay["steps"],
                    "x": x,
                    "y": y,
                    "width": min(window[0], bay["width"]),
                    "height": min(window[1], bay["height"]),
                }
            offset += bay["steps"]
        raise DemoSourceError("step outside the walk")

    def _position(self, local_step: int, image_width: int, image_height: int, window: Tuple[int, int]) -> Tuple[int, int]:
        """Serpentine: pan a row, step up, pan back."""
        win_w = min(window[0], image_width)
        win_h = min(window[1], image_height)
        bands = self._bands(image_height, win_h)
        per_band = self._pan_steps(image_width, win_w)
        cycle = per_band + TRANSITION_STEPS
        band = min(bands - 1, local_step // cycle)
        within = local_step - band * cycle
        stride_y = 0 if bands == 1 else max(1, int(win_h * (1.0 - ROW_OVERLAP)))
        # Ascending: band 0 is the bottom of the shelf.
        y_bottom = min(image_height - win_h, band * stride_y)
        y = max(0, image_height - win_h - y_bottom)
        stride_x = max(1, int(win_w * self.pan_overlap))
        max_x = max(0, image_width - win_w)
        if within < per_band:
            pan = within
            x = min(max_x, pan * stride_x)
        else:
            # The step-up beat: hold x at the end of the row and climb.
            transition = min(TRANSITION_STEPS, max(1, within - per_band + 1))
            pan = per_band - 1
            x = min(max_x, pan * stride_x)
            next_band = min(bands - 1, band + 1)
            next_y_bottom = min(image_height - win_h, next_band * stride_y)
            ratio = transition / float(TRANSITION_STEPS)
            y = int(round(y - (next_y_bottom - y_bottom) * ratio))
        if band % 2 == 1 and within < per_band:
            # Odd bands pan back the other way, like a person reversing direction.
            x = max(0, max_x - x)
        return int(max(0, min(max_x, x))), int(max(0, min(image_height - win_h, y)))

    def frame(
        self,
        sample_ids: Sequence[str],
        step: int,
        window: Optional[Tuple[int, int]] = None,
        quality: int = DEFAULT_QUALITY,
    ) -> Tuple[bytes, Dict[str, Any]]:
        """The JPEG frame the virtual camera sees at this step, plus its geometry."""
        import cv2

        position = self.window_at(sample_ids, step, window)
        image = self._image(position["sample"])
        x, y = position["x"], position["y"]
        w, h = position["width"], position["height"]
        crop = image[y:y + h, x:x + w]
        if crop.size == 0:
            raise DemoSourceError("empty crop")
        ok, buffer = cv2.imencode(".jpg", crop, [cv2.IMWRITE_JPEG_QUALITY, int(quality)])
        if not ok:
            raise DemoSourceError("JPEG encode failed")
        return bytes(buffer), position

    def warm(self, sample_ids: Sequence[str]) -> None:
        """Decode the bays up front so the first demo frame is not the slow one."""
        for sample_id in sample_ids:
            self._size(sample_id)


if __name__ == "__main__":  # smoke: python code/demo.py [sample ...]
    import sys

    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from app import get_available_shelf_images

    samples = {item["id"]: item for item in get_available_shelf_images("add")}
    camera = DemoCamera(samples)
    ids = sys.argv[1:] or ["image6"]
    plan = camera.plan(ids)
    print(f"walk over {plan['samples']}: {plan['steps']} steps, window {plan['window']}")
    started = time.perf_counter()
    for step in range(min(plan["steps"], 24)):
        payload, position = camera.frame(ids, step)
        print(f"  step {step:3d} {position['sample']} x={position['x']:4d} y={position['y']:4d} "
              f"{len(payload) // 1024} KB")
    elapsed = (time.perf_counter() - started) * 1000 / max(1, min(plan["steps"], 24))
    print(f"{elapsed:.1f} ms per frame")