"""WebSocket streaming endpoints: preview, api-input, test."""

import asyncio
import json as _json
import time as _time
import uuid as _uuid

import numpy as np
from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect

from ledgrab.api.dependencies import (
    get_color_strip_store,
    get_processor_manager,
)
from ledgrab.core.processing.processor_manager import ProcessorManager
from ledgrab.storage.color_strip_source import (
    AdvancedPictureColorStripSource,
    ApiInputColorStripSource,
    CompositeColorStripSource,
    NotificationColorStripSource,
    PictureColorStripSource,
)
from ledgrab.storage.color_strip_store import ColorStripStore
from ledgrab.utils import get_logger

logger = get_logger(__name__)

router = APIRouter()


_PREVIEW_ALLOWED_TYPES = {
    "static",
    "gradient",
    "color_cycle",
    "effect",
    "daylight",
    "candlelight",
    "notification",
}


@router.websocket("/api/v1/color-strip-sources/preview/ws")
async def preview_color_strip_ws(
    websocket: WebSocket,
    led_count: int = Query(100),
    fps: int = Query(20),
):
    """Transient preview WebSocket — stream frames for an ad-hoc source config.

    Auth via first-message handshake: ``{"type":"auth","token":"..."}``
    within 3 s of connect.  Use ``?led_count=100&fps=20`` for stream params.

    After auth, waits for a text message containing the full source config
    JSON (must include ``source_type``).  Responds with a JSON metadata message,
    then streams binary RGB frames at the requested FPS.

    Subsequent text messages are treated as config updates: if the source_type
    changed the old stream is replaced; otherwise ``update_source()`` is used.
    """
    from ledgrab.api.auth import accept_and_authenticate_ws

    if await accept_and_authenticate_ws(websocket) is None:
        return

    led_count = max(1, min(1000, led_count))
    fps = max(1, min(60, fps))
    frame_interval = 1.0 / fps

    stream = None
    clock_id = None
    current_source_type = None

    # Helpers ────────────────────────────────────────────────────────────

    def _get_sync_clock_manager():
        """Return the SyncClockManager if available."""
        try:
            mgr = get_processor_manager()
            return getattr(mgr, "_sync_clock_manager", None)
        except Exception as e:
            logger.debug("SyncClockManager not available: %s", e)
            return None

    def _build_source(config: dict):
        """Build a ColorStripSource from a raw config dict, injecting synthetic id/name."""
        from ledgrab.storage.color_strip_source import ColorStripSource

        config.setdefault("id", "__preview__")
        config.setdefault("name", "__preview__")
        return ColorStripSource.from_dict(config)

    def _create_stream(source):
        """Instantiate and start the appropriate stream class for *source*."""
        from ledgrab.core.processing.color_strip_stream_manager import _SIMPLE_STREAM_MAP

        stream_cls = _SIMPLE_STREAM_MAP.get(source.source_type)
        if not stream_cls:
            raise ValueError(f"Unsupported preview source_type: {source.source_type}")
        s = stream_cls(source)
        # Inject gradient store for palette resolution
        if hasattr(s, "set_gradient_store"):
            try:
                from ledgrab.api.dependencies import get_gradient_store

                s.set_gradient_store(get_gradient_store())
            except Exception as e:
                logger.exception("unexpected in gradient_store injection: %s", e)
        if hasattr(s, "configure"):
            s.configure(led_count)
        # Inject sync clock if requested
        cid = getattr(source, "clock_id", None)
        if cid and hasattr(s, "set_clock"):
            scm = _get_sync_clock_manager()
            if scm:
                try:
                    clock_rt = scm.acquire(cid)
                    s.set_clock(clock_rt)
                except Exception as e:
                    logger.warning(f"Preview: could not acquire clock {cid}: {e}")
                    cid = None
            else:
                cid = None
        else:
            cid = None
        s.start()
        return s, cid

    def _stop_stream(s, cid):
        """Stop a stream and release its clock."""
        try:
            s.stop()
        except Exception as e:
            logger.exception("unexpected in _stop_stream s.stop: %s", e)
        if cid:
            scm = _get_sync_clock_manager()
            if scm:
                try:
                    scm.release(cid)
                except Exception as e:
                    logger.exception("unexpected in _stop_stream clock release: %s", e)

    async def _send_meta(source_type: str):
        meta = {"type": "meta", "led_count": led_count, "source_type": source_type}
        await websocket.send_text(_json.dumps(meta))

    # Wait for initial config ────────────────────────────────────────────

    try:
        initial_text = await websocket.receive_text()
    except WebSocketDisconnect:
        return
    except (RuntimeError, ConnectionError) as e:
        logger.debug("ws closed during initial config: %s", e)
        return

    try:
        config = _json.loads(initial_text)
        source_type = config.get("source_type")
        if source_type not in _PREVIEW_ALLOWED_TYPES:
            await websocket.send_text(
                _json.dumps(
                    {
                        "type": "error",
                        "detail": f"source_type must be one of {sorted(_PREVIEW_ALLOWED_TYPES)}",
                    }
                )
            )
            await websocket.close(code=4003, reason="Invalid source_type")
            return
        source = _build_source(config)
        stream, clock_id = _create_stream(source)
        current_source_type = source_type
    except Exception as e:
        logger.error(f"Preview WS: bad initial config: {e}")
        await websocket.send_text(_json.dumps({"type": "error", "detail": str(e)}))
        await websocket.close(code=4003, reason=str(e))
        return

    await _send_meta(current_source_type)
    logger.info(
        f"Preview WS connected: source_type={current_source_type}, led_count={led_count}, fps={fps}"
    )

    # Frame loop ─────────────────────────────────────────────────────────

    try:
        while True:
            # Non-blocking check for incoming config updates
            try:
                msg = await asyncio.wait_for(websocket.receive_text(), timeout=frame_interval)
            except asyncio.TimeoutError:
                msg = None
            except WebSocketDisconnect:
                break

            if msg is not None:
                try:
                    new_config = _json.loads(msg)

                    # Handle "fire" command for notification streams
                    if new_config.get("action") == "fire":
                        from ledgrab.core.processing.notification_stream import (
                            NotificationColorStripStream,
                        )

                        if isinstance(stream, NotificationColorStripStream):
                            stream.fire(
                                app_name=new_config.get("app", ""),
                                color_override=new_config.get("color"),
                            )
                        continue

                    new_type = new_config.get("source_type")
                    if new_type not in _PREVIEW_ALLOWED_TYPES:
                        await websocket.send_text(
                            _json.dumps(
                                {
                                    "type": "error",
                                    "detail": f"source_type must be one of {sorted(_PREVIEW_ALLOWED_TYPES)}",
                                }
                            )
                        )
                        continue
                    new_source = _build_source(new_config)
                    if new_type != current_source_type:
                        # Source type changed — recreate stream
                        _stop_stream(stream, clock_id)
                        stream, clock_id = _create_stream(new_source)
                        current_source_type = new_type
                    else:
                        stream.update_source(new_source)
                        if hasattr(stream, "configure"):
                            stream.configure(led_count)
                    await _send_meta(current_source_type)
                except Exception as e:
                    logger.warning(f"Preview WS: bad config update: {e}")
                    await websocket.send_text(_json.dumps({"type": "error", "detail": str(e)}))

            # Send frame
            colors = stream.get_latest_colors()
            if colors is not None:
                await websocket.send_bytes(colors.tobytes())
            else:
                # Stream hasn't produced a frame yet — send black
                await websocket.send_bytes(b"\x00" * led_count * 3)

    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error(f"Preview WS error: {e}")
    finally:
        if stream is not None:
            _stop_stream(stream, clock_id)
        logger.info("Preview WS disconnected")


@router.websocket("/api/v1/color-strip-sources/{source_id}/ws")
async def css_api_input_ws(
    websocket: WebSocket,
    source_id: str,
):
    """WebSocket for pushing raw LED colors to an api_input source.

    Auth via first-message handshake: ``{"type":"auth","token":"..."}``
    within 3 s of connect.  After auth_ok, accepts JSON frames
    (``{"colors": [[R,G,B], ...]}``) or binary frames (raw RGBRGB... bytes,
    3 bytes per LED).
    """
    from ledgrab.api.auth import accept_and_authenticate_ws

    if await accept_and_authenticate_ws(websocket) is None:
        return

    # Validate source exists and is api_input type
    manager = get_processor_manager()
    try:
        store = get_color_strip_store()
        source = store.get_source(source_id)
    except (ValueError, RuntimeError):
        await websocket.close(code=4004, reason="Source not found")
        return

    if not isinstance(source, ApiInputColorStripSource):
        await websocket.close(code=4003, reason="Source is not api_input type")
        return

    logger.info(f"API input WebSocket connected for source {source_id}")

    try:
        while True:
            message = await websocket.receive()

            if message.get("type") == "websocket.disconnect":
                break

            if "text" in message:
                # JSON frame: {"colors": [[R,G,B], ...]} or {"segments": [...]}
                import json

                try:
                    data = json.loads(message["text"])
                except (json.JSONDecodeError, ValueError) as e:
                    await websocket.send_json({"error": str(e)})
                    continue

                if "segments" in data:
                    # Segment-based path — validate and push
                    try:
                        from ledgrab.api.schemas.color_strip_sources import SegmentPayload

                        seg_dicts = [SegmentPayload(**s).model_dump() for s in data["segments"]]
                    except Exception as e:
                        await websocket.send_json({"error": f"Invalid segment: {e}"})
                        continue
                    streams = manager.color_strip_stream_manager.get_streams_by_source_id(source_id)
                    for stream in streams:
                        if hasattr(stream, "push_segments"):
                            stream.push_segments(seg_dicts)
                    continue

                elif "colors" in data:
                    try:
                        raw_colors = data["colors"]
                        colors_array = np.array(raw_colors, dtype=np.uint8)
                        if colors_array.ndim != 2 or colors_array.shape[1] != 3:
                            await websocket.send_json({"error": "Colors must be [[R,G,B], ...]"})
                            continue
                    except (ValueError, TypeError) as e:
                        await websocket.send_json({"error": str(e)})
                        continue
                else:
                    await websocket.send_json(
                        {"error": "JSON frame must contain 'colors' or 'segments'"}
                    )
                    continue

            elif "bytes" in message:
                # Binary frame: raw RGBRGB... bytes (3 bytes per LED)
                raw_bytes = message["bytes"]
                if len(raw_bytes) % 3 != 0:
                    await websocket.send_json({"error": "Binary data must be multiple of 3 bytes"})
                    continue
                colors_array = np.frombuffer(raw_bytes, dtype=np.uint8).reshape(-1, 3)

            else:
                continue

            # Push to all running streams (colors_array path only reaches here)
            streams = manager.color_strip_stream_manager.get_streams_by_source_id(source_id)
            for stream in streams:
                if hasattr(stream, "push_colors"):
                    stream.push_colors(colors_array)

    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error(f"API input WebSocket error for source {source_id}: {e}")
    finally:
        logger.info(f"API input WebSocket disconnected for source {source_id}")


@router.websocket("/api/v1/color-strip-sources/{source_id}/test/ws")
async def test_color_strip_ws(
    websocket: WebSocket,
    source_id: str,
    led_count: int = Query(100),
    fps: int = Query(20),
):
    """WebSocket for real-time CSS source preview.

    Auth via first-message handshake: ``{"type":"auth","token":"..."}``
    within 3 s of connect.  After auth_ok, server sends JSON metadata
    (source_type, led_count, calibration segments), then binary RGB frames
    (``led_count * 3`` bytes).
    """
    from ledgrab.api.auth import accept_and_authenticate_ws

    if await accept_and_authenticate_ws(websocket) is None:
        return

    # Validate source exists
    store: ColorStripStore = get_color_strip_store()
    try:
        source = store.get_source(source_id)
    except ValueError as e:
        await websocket.close(code=4004, reason=str(e))
        return

    # Acquire stream – unique consumer ID per WS to avoid release races
    manager: ProcessorManager = get_processor_manager()
    csm = manager.color_strip_stream_manager
    consumer_id = f"__test_{_uuid.uuid4().hex[:8]}__"
    try:
        stream = csm.acquire(source_id, consumer_id)
    except Exception as e:
        logger.error(f"CSS test: failed to acquire stream for {source_id}: {e}")
        await websocket.close(code=4003, reason=str(e))
        return

    # Configure LED count for auto-sizing streams
    if hasattr(stream, "configure"):
        stream.configure(max(1, led_count))

    # Reject picture sources with 0 calibration LEDs (no edges configured)
    if stream.led_count <= 0:
        csm.release(source_id, consumer_id)
        await websocket.close(
            code=4005,
            reason="No LEDs configured. Open Calibration and set LED counts for each edge.",
        )
        return

    # Clamp FPS to sane range
    fps = max(1, min(60, fps))
    _frame_interval = 1.0 / fps

    logger.info(f"CSS test WebSocket connected for {source_id} (fps={fps})")

    try:
        from ledgrab.core.processing.composite_stream import CompositeColorStripStream

        from ledgrab.core.processing.api_input_stream import ApiInputColorStripStream

        is_api_input = isinstance(stream, ApiInputColorStripStream)
        _last_push_gen = 0  # track api_input push generation to skip unchanged frames

        # Send metadata as first message
        is_picture = isinstance(source, (PictureColorStripSource, AdvancedPictureColorStripSource))
        is_composite = isinstance(source, CompositeColorStripSource)
        meta: dict = {
            "type": "meta",
            "source_type": source.source_type,
            "source_name": source.name,
            "led_count": stream.led_count,
        }
        if is_picture and stream.calibration:
            cal = stream.calibration
            total = cal.get_total_leds()
            offset = cal.offset % total if total > 0 else 0
            edges = []
            for seg in cal.segments:
                # Compute output indices matching PixelMapper logic
                indices = list(range(seg.led_start, seg.led_start + seg.led_count))
                if seg.reverse:
                    indices = indices[::-1]
                if offset > 0:
                    indices = [(idx + offset) % total for idx in indices]
                edges.append({"edge": seg.edge, "indices": indices})
            meta["edges"] = edges
            meta["border_width"] = cal.border_width
        if is_composite and hasattr(source, "layers"):
            # Send layer info for composite preview
            enabled_layers = [layer for layer in source.layers if layer.get("enabled", True)]
            layer_infos = []  # [{name, id, is_notification, has_brightness, ...}, ...]
            for layer in enabled_layers:
                info = {
                    "id": layer["source_id"],
                    "name": layer.get("source_id", "?"),
                    "is_notification": False,
                    "has_brightness": bool(layer.get("brightness_source_id")),
                }
                try:
                    layer_src = store.get_source(layer["source_id"])
                    info["name"] = layer_src.name
                    info["is_notification"] = isinstance(layer_src, NotificationColorStripSource)
                    if isinstance(
                        layer_src, (PictureColorStripSource, AdvancedPictureColorStripSource)
                    ):
                        info["is_picture"] = True
                        if hasattr(layer_src, "calibration") and layer_src.calibration:
                            info["calibration_led_count"] = layer_src.calibration.get_total_leds()
                except (ValueError, KeyError):
                    pass
                layer_infos.append(info)
            meta["layers"] = [li["name"] for li in layer_infos]
            meta["layer_infos"] = layer_infos
        await websocket.send_text(_json.dumps(meta))

        # For api_input: send the current buffer immediately so the client
        # gets a frame right away (fallback color if inactive) rather than
        # leaving the canvas blank/stale until external data arrives.
        if is_api_input:
            initial_colors = stream.get_latest_colors()
            if initial_colors is not None:
                await websocket.send_bytes(initial_colors.tobytes())

        # For picture sources, grab the live stream for frame preview
        _frame_live = None
        if is_picture and hasattr(stream, "live_stream"):
            _frame_live = stream.live_stream
        _last_aux_time = 0.0
        _AUX_INTERVAL = 0.08  # send JPEG preview / brightness updates ~12 FPS
        _frame_dims_sent = False  # send frame dimensions once with first JPEG

        # Stream binary RGB frames at ~20 Hz
        while True:
            # For composite sources, send per-layer data like target preview does
            if is_composite and isinstance(stream, CompositeColorStripStream):
                layer_colors = stream.get_layer_colors()
                composite_colors = stream.get_latest_colors()
                if composite_colors is not None and layer_colors and len(layer_colors) > 1:
                    led_count = composite_colors.shape[0]
                    rgb_size = led_count * 3
                    # Wire format: [0xFE] [layer_count] [led_count_hi] [led_count_lo] [layer0_rgb...] ... [composite_rgb]
                    header = bytes(
                        [0xFE, len(layer_colors), (led_count >> 8) & 0xFF, led_count & 0xFF]
                    )
                    parts = [header]
                    for lc in layer_colors:
                        if lc is not None and lc.shape[0] == led_count:
                            parts.append(lc.tobytes())
                        else:
                            parts.append(b"\x00" * rgb_size)
                    parts.append(composite_colors.tobytes())
                    await websocket.send_bytes(b"".join(parts))
                elif composite_colors is not None:
                    await websocket.send_bytes(composite_colors.tobytes())
            else:
                # For api_input: only send when new data was pushed
                if is_api_input:
                    gen = stream.push_generation
                    if gen != _last_push_gen:
                        _last_push_gen = gen
                        colors = stream.get_latest_colors()
                        if colors is not None:
                            await websocket.send_bytes(colors.tobytes())
                else:
                    colors = stream.get_latest_colors()
                    if colors is not None:
                        await websocket.send_bytes(colors.tobytes())

            # Periodically send auxiliary data (frame preview, brightness)
            now = _time.monotonic()
            if now - _last_aux_time >= _AUX_INTERVAL:
                _last_aux_time = now

                # Send brightness values for composite layers
                if is_composite and isinstance(stream, CompositeColorStripStream):
                    try:
                        bri_values = stream.get_layer_brightness()
                        if any(v is not None for v in bri_values):
                            bri_msg = {
                                "type": "brightness",
                                "values": [
                                    round(v * 100) if v is not None else None for v in bri_values
                                ],
                            }
                            await websocket.send_text(_json.dumps(bri_msg))
                    except Exception as e:
                        logger.exception("unexpected in brightness send: %s", e)

                # Send JPEG frame preview for picture sources
                if _frame_live:
                    try:
                        frame = _frame_live.get_latest_frame()
                        if frame is not None and frame.image is not None:
                            from ledgrab.utils.image_codec import encode_jpeg, resize_image

                            img = frame.image
                            # Ensure 3-channel RGB (some engines may produce BGRA)
                            if img.ndim == 3 and img.shape[2] == 4:
                                img = img[:, :, :3]
                            h, w = img.shape[:2]
                            # Send frame dimensions once so client can compute border overlay
                            if not _frame_dims_sent:
                                _frame_dims_sent = True
                                await websocket.send_text(
                                    _json.dumps(
                                        {
                                            "type": "frame_dims",
                                            "width": w,
                                            "height": h,
                                        }
                                    )
                                )
                            # Downscale for bandwidth
                            scale = min(960 / w, 540 / h, 1.0)
                            if scale < 1.0:
                                new_w = max(1, int(w * scale))
                                new_h = max(1, int(h * scale))
                                img = resize_image(img, new_w, new_h)
                            # Wire format: [0xFD] [jpeg_bytes]
                            await websocket.send_bytes(b"\xfd" + encode_jpeg(img, quality=70))
                    except Exception as e:
                        logger.warning(f"JPEG frame preview error: {e}")

            await asyncio.sleep(_frame_interval)
    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error(f"CSS test WebSocket error for {source_id}: {e}")
    finally:
        csm.release(source_id, consumer_id)
        logger.info(f"CSS test WebSocket disconnected for {source_id}")
