#!/usr/bin/env python3
"""Headless 3-display LiveKit wall via GStreamer kmssink.

  1. grid        — 5x4 mosaic of cam-01..cam-20
  2. speaker     — designated speaker camera (DISPLAY_SPEAKER_CAMERA)
  3. screenshare — STUB until the production LiveKit server is online

    .venv/bin/python run_displays.py --list
    .venv/bin/python run_displays.py --test-pattern
    .venv/bin/python run_displays.py
"""
from __future__ import annotations

import argparse
import asyncio
import logging
import signal
import sys
import time
from typing import Optional
from urllib.parse import urlparse

import numpy as np

from config import load_config
from display_grid import (
    CameraGrid,
    ROLES,
    bgra_from_livekit,
    camera_ids,
    is_camera_identity,
    placeholder,
)
from drm_outputs import DrmOutput, format_outputs, list_outputs, select_outputs
from gst_sink import KmsPipeline
from tokens import make_viewer_token

log = logging.getLogger("cameras.display")

_PATTERNS = ("smpte", "ball", "snow")
_PLACE_W, _PLACE_H = 1280, 720


def _public_url(cfg) -> str:
    if cfg.livekit_public_url:
        return cfg.livekit_public_url.rstrip("/")
    parsed = urlparse(cfg.url)
    host = parsed.hostname or "127.0.0.1"
    if host in ("127.0.0.1", "localhost"):
        host = "10.200.0.20"
    scheme = parsed.scheme or "ws"
    port = parsed.port
    if port:
        return f"{scheme}://{host}:{port}"
    return f"{scheme}://{host}"


def _bind_roles(roles: list[str], connectors: list[str]) -> list[tuple[str, DrmOutput]]:
    roles = [r.strip().lower() for r in roles if r.strip()]
    for r in roles:
        if r not in ROLES:
            raise KeyError(f"unknown display role {r!r}; expected {ROLES}")
    names = connectors if connectors else None
    outs = select_outputs(n=len(roles), names=names)
    if len(outs) < len(roles):
        log.warning("wanted %d displays, found %d DRM connectors", len(roles), len(outs))
    slots = list(zip(roles, outs))
    for role, out in slots:
        log.info("role %s -> %s (%s connector-id=%s %s)",
                 role, out.name, out.status, out.connector_id, out.driver)
    return slots


class Wall:
    def __init__(self, cfg, slots: list[tuple[str, DrmOutput]]):
        self.cfg = cfg
        self.slots = {role: KmsPipeline(out, role) for role, out in slots}
        prefix = cfg.participant_prefix or "cam"
        n = cfg.display_grid_cols * cfg.display_grid_rows
        self.grid = CameraGrid(
            width=cfg.display_grid_width,
            height=cfg.display_grid_height,
            cols=cfg.display_grid_cols,
            rows=cfg.display_grid_rows,
            identities=camera_ids(n, prefix),
        )
        self.speaker_camera = (cfg.display_speaker_camera or "cam-01").strip()
        self.active_speaker: Optional[str] = None
        self._speaker_frame: Optional[np.ndarray] = None
        self._tasks: dict[str, asyncio.Task] = {}
        self._stop = asyncio.Event()
        # Screenshare is a stub until production LiveKit is online.
        self._share_placeholder = placeholder(_PLACE_W, _PLACE_H, [
            "Screen share",
            "stub — waiting for production LiveKit",
            _public_url(cfg),
        ])
        cam = self.speaker_camera
        self._speaker_placeholder = placeholder(_PLACE_W, _PLACE_H, [
            "Speaker camera",
            cam if cam.lower() != "active" else "(following active speaker)",
        ])

    def speaker_identity(self) -> Optional[str]:
        if self.speaker_camera.lower() == "active":
            return self.active_speaker
        return self.speaker_camera

    def _push_sink(self, role: str, data: bytes, width: int, height: int, fps: int) -> None:
        sink = self.slots.get(role)
        if sink is None or not sink.out.connected:
            return
        if not sink.alive or sink.width != width or sink.height != height:
            sink.start_appsrc(width, height, fps)
        if not sink.push(data):
            sink.start_appsrc(width, height, fps)
            sink.push(data)
        err = sink.gst_errors()
        if err:
            for line in err.strip().splitlines()[-6:]:
                log.warning("%s gst: %s", role, line)

    async def _output_loops(self) -> None:
        grid_fps = max(1, int(self.cfg.display_grid_fps))
        grid_period = 1.0 / grid_fps
        spk_period = 1.0 / max(1, int(self.cfg.fps))
        next_grid = time.monotonic()
        next_spk = time.monotonic()
        while not self._stop.is_set():
            now = time.monotonic()
            if now >= next_grid:
                self._push_sink(
                    "grid", self.grid.render(),
                    self.grid.width, self.grid.height, grid_fps,
                )
                next_grid = now + grid_period
            if now >= next_spk:
                spk = self._speaker_frame
                if spk is None:
                    img = self._speaker_placeholder
                else:
                    img = spk
                self._push_sink("speaker", img.tobytes(), img.shape[1], img.shape[0], self.cfg.fps)
                share = self._share_placeholder
                self._push_sink("screenshare", share.tobytes(), share.shape[1], share.shape[0], self.cfg.fps)
                next_spk = now + spk_period
            await asyncio.sleep(0.005)

    async def _pump_camera(self, track, identity: str) -> None:
        import livekit.rtc as rtc

        stream = rtc.VideoStream(track)
        try:
            async for event in stream:
                frame = event.frame
                if frame.width <= 0 or frame.height <= 0:
                    continue
                try:
                    converted = frame.convert(rtc.VideoBufferType.BGRA)
                except Exception as exc:
                    log.warning("%s: convert failed: %s", identity, exc)
                    continue
                bgra = bgra_from_livekit(converted.data, converted.width, converted.height)
                self.grid.set_frame(identity, bgra)
                if identity == self.speaker_identity():
                    self._speaker_frame = bgra
        finally:
            await stream.aclose()
            self.grid.clear(identity)
            if identity == self.speaker_identity():
                self._speaker_frame = None

    def _track_key(self, participant_identity: str, track) -> str:
        return f"{participant_identity}:{getattr(track, 'sid', id(track))}"

    def attach(self, track, publication, participant, rtc) -> None:
        ident = participant.identity
        if ident == self.cfg.display_identity:
            return
        key = self._track_key(ident, track)
        old = self._tasks.pop(key, None)
        if old:
            old.cancel()
        # Screenshare subscribe/render is stubbed until production LiveKit.
        if is_camera_identity(ident, self.cfg.participant_prefix):
            self._tasks[key] = asyncio.create_task(
                self._pump_camera(track, ident), name=f"cam-{ident}")

    def subscribe_if_wanted(self, participant, rtc) -> None:
        ident = participant.identity
        if ident == self.cfg.display_identity:
            return
        cam = is_camera_identity(ident, self.cfg.participant_prefix)
        for pub in participant.track_publications.values():
            if pub.kind != rtc.TrackKind.KIND_VIDEO:
                continue
            want = cam
            if want and not pub.subscribed:
                log.info("subscribe %s from %s (source=%s)", pub.sid, ident, pub.source)
                pub.set_subscribed(True)

    def detach(self, track, participant) -> None:
        key = self._track_key(participant.identity, track)
        task = self._tasks.pop(key, None)
        if task:
            task.cancel()

    async def run(self) -> int:
        import livekit.rtc as rtc

        room = rtc.Room()

        @room.on("track_subscribed")
        def _on_sub(track, publication, participant):
            if track.kind == rtc.TrackKind.KIND_VIDEO:
                log.info("subscribed %s from %s", publication.sid, participant.identity)
                self.attach(track, publication, participant, rtc)

        @room.on("track_unsubscribed")
        def _on_unsub(track, publication, participant):
            self.detach(track, participant)

        @room.on("track_published")
        def _on_pub(publication, participant):
            self.subscribe_if_wanted(participant, rtc)

        @room.on("participant_connected")
        def _on_pc(participant):
            log.info("participant joined: %s", participant.identity)
            self.subscribe_if_wanted(participant, rtc)

        @room.on("active_speakers_changed")
        def _on_spk(speakers):
            cams = [
                p.identity for p in speakers
                if is_camera_identity(p.identity, self.cfg.participant_prefix)
            ]
            self.active_speaker = cams[0] if cams else None

        token = make_viewer_token(
            self.cfg.api_key, self.cfg.api_secret, self.cfg.room, self.cfg.display_identity)
        log.info("connecting to %s room=%s as %s", self.cfg.url, self.cfg.room, self.cfg.display_identity)
        await room.connect(self.cfg.url, token, rtc.RoomOptions(auto_subscribe=False))

        for rp in room.remote_participants.values():
            self.subscribe_if_wanted(rp, rtc)
            for pub in rp.track_publications.values():
                if pub.track and pub.track.kind == rtc.TrackKind.KIND_VIDEO:
                    self.attach(pub.track, pub, rp, rtc)

        log.info(
            "wall up: grid=%dx%d speaker_camera=%s screenshare=STUB",
            self.grid.cols, self.grid.rows, self.speaker_camera,
        )
        loops = asyncio.create_task(self._output_loops(), name="output-loops")
        loop = asyncio.get_running_loop()
        for sig in (signal.SIGINT, signal.SIGTERM):
            try:
                loop.add_signal_handler(sig, self._stop.set)
            except NotImplementedError:
                pass
        try:
            await self._stop.wait()
        finally:
            loops.cancel()
            for t in self._tasks.values():
                t.cancel()
            for sink in self.slots.values():
                sink.stop()
            await room.disconnect()
        return 0


def run_test_pattern(slots: list[tuple[str, DrmOutput]]) -> int:
    sinks = []
    for i, (role, out) in enumerate(slots):
        sink = KmsPipeline(out, role)
        sink.start_test_pattern(_PATTERNS[i % len(_PATTERNS)])
        sinks.append(sink)
    log.info("test patterns on %d sink(s); Ctrl-C to stop", sum(1 for s in sinks if s.alive))
    stop = False

    def _stop(*_):
        nonlocal stop
        stop = True

    signal.signal(signal.SIGINT, _stop)
    signal.signal(signal.SIGTERM, _stop)
    try:
        while not stop:
            for s in sinks:
                err = s.gst_errors()
                if err:
                    for line in err.strip().splitlines()[-8:]:
                        log.warning("%s gst: %s", s.identity, line)
            if any(s.out.connected and not s.alive and s.mode == "dead" for s in sinks):
                return 1
            time.sleep(0.5)
    finally:
        for s in sinks:
            s.stop()
    return 0


def main(argv: Optional[list[str]] = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--list", action="store_true", help="list DRM connectors and exit")
    ap.add_argument("--test-pattern", action="store_true",
                    help="drive displays with GStreamer videotestsrc (no LiveKit)")
    ap.add_argument("--connectors", default="",
                    help="comma-separated DRM names, e.g. HDMI-A-7,HDMI-A-5,HDMI-A-6")
    ap.add_argument("--roles", default="",
                    help="comma-separated roles (default grid,speaker,screenshare)")
    ap.add_argument("--speaker-camera", default="",
                    help="cam-NN or 'active' (overrides DISPLAY_SPEAKER_CAMERA)")
    args = ap.parse_args(argv)

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [display] %(levelname)s %(message)s",
    )

    if args.list:
        print(format_outputs(list_outputs()) or "(no DRM connectors)")
        return 0

    cfg = load_config()
    logging.getLogger().setLevel(getattr(logging, cfg.log_level, logging.INFO))
    if args.speaker_camera:
        cfg.display_speaker_camera = args.speaker_camera
    roles = [p.strip() for p in args.roles.split(",") if p.strip()] or cfg.display_roles
    connectors = [p.strip() for p in args.connectors.split(",") if p.strip()] or cfg.display_connectors
    try:
        slots = _bind_roles(roles, connectors)
    except KeyError as exc:
        log.error("%s", exc)
        return 2
    if not slots:
        log.error("no DRM outputs to drive")
        return 2

    if args.test_pattern:
        return run_test_pattern(slots)
    return asyncio.run(Wall(cfg, slots).run())


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