#!/usr/bin/env python3
"""JPEG proxy: prefetch from Frigate per camera/size and serve cover-cropped JPEGs."""

from __future__ import annotations

import hashlib
import io
import json
import os
import re
import threading
import time
from dataclasses import dataclass, field
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse

import requests
from flask import Flask, abort, request
from PIL import Image

JPEG_QUALITY = int(os.environ.get("JPEG_QUALITY", "40"))
FETCH_TIMEOUT = float(os.environ.get("FETCH_TIMEOUT", "5"))
PREFETCH_HZ = float(os.environ.get("PREFETCH_HZ", "10"))
DEFAULT_SIZE = int(os.environ.get("DEFAULT_SIZE", "480"))
MAX_SIZE = int(os.environ.get("MAX_SIZE", "720"))
MIN_SIZE = int(os.environ.get("MIN_SIZE", "64"))
MAX_ZOOM = float(os.environ.get("MAX_ZOOM", "3.0"))
MIN_ZOOM = float(os.environ.get("MIN_ZOOM", "1.0"))

_NAME_RE = re.compile(r"^[a-z0-9][a-z0-9_-]{0,31}$")


@dataclass(frozen=True)
class CameraSource:
    name: str
    url: str
    zoom: float = 1.0


@dataclass
class StreamConfig:
    source: CameraSource
    width: int
    height: int
    zoom: float

    @property
    def key(self) -> tuple[str, int, int, float]:
        return (self.source.name, self.width, self.height, round(self.zoom, 3))

    @property
    def label(self) -> str:
        if self.zoom != 1.0:
            return f"{self.source.name}:{self.width}x{self.height}@{self.zoom:g}x"
        return f"{self.source.name}:{self.width}x{self.height}"


@dataclass
class CameraStream:
    config: StreamConfig
    lock: threading.Lock = field(default_factory=threading.Lock)
    jpeg: bytes = b""
    etag: str = ""
    width: int = 0
    height: int = 0

    def snapshot(self) -> tuple[bytes, str, int, int]:
        with self.lock:
            return self.jpeg, self.etag, self.width, self.height

    def store(self, body: bytes, etag: str, width: int, height: int) -> None:
        with self.lock:
            self.jpeg = body
            self.etag = etag
            self.width = width
            self.height = height


class StreamRegistry:
    def __init__(self, sources: dict[str, CameraSource]) -> None:
        self._sources = sources
        self._streams: dict[tuple[str, int, int, float], CameraStream] = {}
        self._lock = threading.Lock()
        self._prefetch_started: set[tuple[str, int, int, float]] = set()

    @staticmethod
    def _clamp(value: int) -> int:
        return max(MIN_SIZE, min(MAX_SIZE, value))

    @staticmethod
    def _clamp_zoom(value: float) -> float:
        return max(MIN_ZOOM, min(MAX_ZOOM, value))

    @classmethod
    def parse_zoom(cls, args, source: CameraSource) -> float:
        raw = args.get("zoom", args.get("z"))
        if raw is None:
            return cls._clamp_zoom(source.zoom)
        return cls._clamp_zoom(float(raw))

    @classmethod
    def parse_dimensions(cls, args) -> tuple[int, int]:
        raw_w = args.get("w", args.get("width"))
        raw_h = args.get("h", args.get("height"))
        raw_size = args.get("size")

        if raw_size is not None:
            size = cls._clamp(int(raw_size))
            return size, size

        if raw_w is None and raw_h is None:
            size = cls._clamp(DEFAULT_SIZE)
            return size, size

        width = cls._clamp(int(raw_w if raw_w is not None else raw_h))
        height = cls._clamp(int(raw_h if raw_h is not None else raw_w))
        return width, height

    def get(self, camera: str, width: int, height: int, zoom: float) -> CameraStream:
        if camera not in self._sources:
            abort(404)
        key = (camera, width, height, round(zoom, 3))
        with self._lock:
            stream = self._streams.get(key)
            if stream is None:
                cfg = StreamConfig(
                    source=self._sources[camera],
                    width=width,
                    height=height,
                    zoom=zoom,
                )
                stream = CameraStream(config=cfg)
                self._streams[key] = stream
            if key not in self._prefetch_started:
                self._prefetch_started.add(key)
                threading.Thread(
                    target=_prefetch_loop,
                    args=(stream,),
                    name=f"prefetch-{camera}-{width}x{height}-z{zoom:g}",
                    daemon=True,
                ).start()
            return stream

    def prefetch_all(self, sizes: list[int]) -> None:
        for source in self._sources.values():
            for size in sizes:
                self.get(source.name, size, size, self._clamp_zoom(source.zoom))

    def all_streams(self) -> list[CameraStream]:
        with self._lock:
            return list(self._streams.values())


def _cover_crop(
    frame: Image.Image, target_w: int, target_h: int, zoom: float = 1.0
) -> Image.Image:
    scale = max(target_w / frame.width, target_h / frame.height) * zoom
    new_w = max(1, round(frame.width * scale))
    new_h = max(1, round(frame.height * scale))
    if (new_w, new_h) != frame.size:
        frame = frame.resize((new_w, new_h), Image.BILINEAR)
    left = max(0, (new_w - target_w) // 2)
    top = max(0, (new_h - target_h) // 2)
    return frame.crop((left, top, left + target_w, top + target_h))


def _encode(
    frame: Image.Image, target_w: int, target_h: int, zoom: float = 1.0
) -> tuple[bytes, int, int]:
    frame = _cover_crop(frame.convert("RGB"), target_w, target_h, zoom)
    buf = io.BytesIO()
    frame.save(
        buf,
        format="JPEG",
        quality=JPEG_QUALITY,
        subsampling=2,
        optimize=False,
        progressive=False,
    )
    return buf.getvalue(), frame.width, frame.height


def _frigate_url(base: str, width: int, height: int, zoom: float = 1.0) -> str:
    fetch_h = round(max(width, height) * zoom)
    parsed = urlparse(base)
    query = parse_qs(parsed.query)
    query["h"] = [str(fetch_h)]
    return urlunparse(parsed._replace(query=urlencode(query, doseq=True)))


def _parse_sources() -> dict[str, CameraSource]:
    raw = os.environ.get("CAMERAS", "").strip()
    if raw:
        items = json.loads(raw)
        if not isinstance(items, list) or not items:
            raise ValueError("CAMERAS must be a non-empty JSON array")
        sources: dict[str, CameraSource] = {}
        for item in items:
            name = str(item["name"]).lower()
            if not _NAME_RE.match(name):
                raise ValueError(f"invalid camera name: {name!r}")
            if name in sources:
                raise ValueError(f"duplicate camera name: {name!r}")
            url = str(item["url"])
            zoom = float(item.get("zoom", 1.0))
            sources[name] = CameraSource(name=name, url=url, zoom=zoom)
        return sources

    url = os.environ.get(
        "FRIGATE_URL", "http://10.0.1.197:5050/api/doorbell/latest.jpg"
    )
    name = os.environ.get("CAMERA_NAME", "doorbell").lower()
    return {name: CameraSource(name=name, url=url)}


def _parse_prefetch_sizes() -> list[int]:
    raw = os.environ.get("PREFETCH_SIZES", "480,720").strip()
    if not raw:
        return []
    sizes: list[int] = []
    for part in raw.split(","):
        part = part.strip()
        if part:
            sizes.append(StreamRegistry._clamp(int(part)))
    return sizes


def _prefetch_loop(stream: CameraStream) -> None:
    cfg = stream.config
    session = requests.Session()
    interval = 1.0 / max(PREFETCH_HZ, 0.5)
    while True:
        started = time.monotonic()
        try:
            source_url = _frigate_url(cfg.source.url, cfg.width, cfg.height, cfg.zoom)
            response = session.get(source_url, timeout=FETCH_TIMEOUT)
            response.raise_for_status()
            with Image.open(io.BytesIO(response.content)) as image:
                encoded, width, height = _encode(
                    image, cfg.width, cfg.height, cfg.zoom
                )
            etag = '"' + hashlib.md5(encoded).hexdigest() + '"'
            stream.store(encoded, etag, width, height)
        except Exception as exc:  # noqa: BLE001 — keep last good frame
            print(f"[{cfg.label}] prefetch failed: {exc}", flush=True)
        delay = interval - (time.monotonic() - started)
        if delay > 0:
            time.sleep(delay)


def _jpeg_response(stream: CameraStream):
    body, etag, width, height = stream.snapshot()
    if not body:
        return "Image not ready", 503, {"Content-Type": "text/plain"}

    if request.headers.get("If-None-Match") == etag:
        return ("", 304, {"ETag": etag, "Cache-Control": "no-cache"})

    return (
        body,
        200,
        {
            "Content-Type": "image/jpeg",
            "Content-Length": str(len(body)),
            "ETag": etag,
            "Last-Modified": time.strftime("%a, %d %b %Y %H:%M:%S GMT", time.gmtime()),
            "Cache-Control": "no-cache",
            "Accept-Ranges": "none",
            "X-Camera": stream.config.source.name,
            "X-Image-Size": f"{width}x{height}",
            "X-Zoom": f"{stream.config.zoom:g}",
        },
    )


def create_app() -> Flask:
    sources = _parse_sources()
    registry = StreamRegistry(sources)
    primary = next(iter(sources))
    registry.prefetch_all(_parse_prefetch_sizes())

    app = Flask(__name__)

    def _serve(camera: str):
        source = registry._sources[camera]
        width, height = registry.parse_dimensions(request.args)
        zoom = registry.parse_zoom(request.args, source)
        return _jpeg_response(registry.get(camera, width, height, zoom))

    @app.get("/health")
    def health():
        result = {}
        for stream in registry.all_streams():
            body, _, width, height = stream.snapshot()
            label = stream.config.label
            result[label] = {
                "ready": bool(body),
                "bytes": len(body),
                "size": f"{width}x{height}",
                "target": f"{stream.config.width}x{stream.config.height}",
                "zoom": stream.config.zoom,
                "source": stream.config.source.url,
            }
        return {"streams": result}, 200 if any(s["ready"] for s in result.values()) else 503

    @app.get("/")
    @app.get("/latest.jpg")
    def latest_primary():
        return _serve(primary)

    for name in sources:
        handler = (lambda camera=name: lambda: _serve(camera))()
        app.add_url_rule(f"/{name}/", endpoint=f"camera-{name}", view_func=handler)
        app.add_url_rule(
            f"/{name}/latest.jpg",
            endpoint=f"camera-{name}-jpg",
            view_func=handler,
        )

    return app


app = create_app()

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", "80")), threaded=True)