"""Key-colors preview/test endpoints (REST + WebSocket)."""

import asyncio

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

from ledgrab.api.auth import AuthRequired
from ledgrab.api.dependencies import (
    get_color_strip_store,
    get_device_store,
    get_picture_source_store,
    get_pp_template_store,
    get_processor_manager,
    get_template_store,
)
from ledgrab.core.processing.processor_manager import ProcessorManager
from ledgrab.storage import DeviceStore
from ledgrab.storage.color_strip_store import ColorStripStore
from ledgrab.storage.picture_source_store import PictureSourceStore
from ledgrab.storage.template_store import TemplateStore
from ledgrab.utils import get_logger

logger = get_logger(__name__)

router = APIRouter()


@router.post(
    "/api/v1/color-strip-sources/{source_id}/key-colors/test",
    tags=["Color Strip Sources"],
)
async def test_key_colors_source(
    source_id: str,
    _auth: AuthRequired,
    store: ColorStripStore = Depends(get_color_strip_store),
    source_store: PictureSourceStore = Depends(get_picture_source_store),
    template_store: TemplateStore = Depends(get_template_store),
    processor_manager: ProcessorManager = Depends(get_processor_manager),
    device_store: DeviceStore = Depends(get_device_store),
    pp_template_store=Depends(get_pp_template_store),
):
    """Test a key_colors source: capture a frame, extract colors from each rectangle."""
    from ledgrab.storage.color_strip_source import KeyColorsColorStripSource
    from ledgrab.core.capture.screen_capture import (
        calculate_average_color,
        calculate_dominant_color,
        calculate_median_color,
    )
    from ledgrab.core.capture_engines import EngineRegistry
    from ledgrab.core.filters import FilterRegistry, ImagePool
    from ledgrab.storage.picture_source import (
        ScreenCapturePictureSource,
        StaticImagePictureSource,
    )
    from ledgrab.utils.image_codec import encode_jpeg_data_uri

    stream = None
    try:
        source = store.get_source(source_id)
        if not isinstance(source, KeyColorsColorStripSource):
            raise HTTPException(status_code=400, detail="Source is not a key_colors type")

        if not source.rectangles:
            raise HTTPException(status_code=400, detail="No screen regions configured")
        if not source.picture_source_id:
            raise HTTPException(status_code=400, detail="No picture source configured")

        # Resolve picture source and capture a frame
        chain = source_store.resolve_stream_chain(source.picture_source_id)
        raw_stream = chain["raw_stream"]

        from ledgrab.utils.image_codec import load_image_file

        if isinstance(raw_stream, StaticImagePictureSource):
            from ledgrab.api.dependencies import get_asset_store as _get_asset_store

            asset_store = _get_asset_store()
            image_path = (
                asset_store.get_file_path(raw_stream.image_asset_id)
                if raw_stream.image_asset_id
                else None
            )
            if not image_path:
                raise HTTPException(status_code=400, detail="Image asset not found")
            image = load_image_file(image_path)
        elif isinstance(raw_stream, ScreenCapturePictureSource):
            capture_template = template_store.get_template(raw_stream.capture_template_id)
            display_index = raw_stream.display_index

            if capture_template.engine_type not in EngineRegistry.get_available_engines():
                raise HTTPException(
                    status_code=400, detail=f"Engine '{capture_template.engine_type}' not available"
                )

            locked = processor_manager.get_display_lock_info(display_index)
            if locked:
                try:
                    device_name = device_store.get_device(locked).name
                except Exception:
                    device_name = locked
                raise HTTPException(
                    status_code=409,
                    detail=f"Display {display_index} is captured by '{device_name}'. Stop it first.",
                )

            stream = EngineRegistry.create_stream(
                capture_template.engine_type, display_index, capture_template.engine_config
            )
            stream.initialize()
            sc = stream.capture_frame()
            if sc is None:
                raise RuntimeError("No frame captured")
            image = sc.image
        else:
            raise HTTPException(status_code=400, detail="Unsupported picture source type")

        # Apply postprocessing filters
        pp_ids = chain.get("postprocessing_template_ids", [])
        if pp_ids and pp_template_store:
            pool = ImagePool()
            for pp_id in pp_ids:
                try:
                    pp = pp_template_store.get_template(pp_id)
                    for fi in pp_template_store.resolve_filter_instances(pp.filters):
                        f = FilterRegistry.create_instance(fi.filter_id, fi.options)
                        result = f.process_image(image, pool)
                        if result is not None:
                            image = result
                except Exception as e:
                    logger.exception("unexpected in postprocessing pp_id=%s: %s", pp_id, e)

        # Extract colors from each rectangle
        h, w = image.shape[:2]
        calc_fns = {
            "average": calculate_average_color,
            "median": calculate_median_color,
            "dominant": calculate_dominant_color,
        }
        calc_fn = calc_fns.get(source.interpolation_mode, calculate_average_color)

        result_rects = []
        for rect in source.rectangles:
            px_x = max(0, int(rect.x * w))
            px_y = max(0, int(rect.y * h))
            px_w = max(1, int(rect.width * w))
            px_h = max(1, int(rect.height * h))
            px_x, px_y = min(px_x, w - 1), min(px_y, h - 1)
            px_w, px_h = min(px_w, w - px_x), min(px_h, h - px_y)
            sub_img = image[px_y : px_y + px_h, px_x : px_x + px_w]
            r, g, b = calc_fn(sub_img)
            result_rects.append(
                {
                    "name": rect.name,
                    "x": rect.x,
                    "y": rect.y,
                    "width": rect.width,
                    "height": rect.height,
                    "color": {
                        "r": int(r),
                        "g": int(g),
                        "b": int(b),
                        "hex": f"#{int(r):02x}{int(g):02x}{int(b):02x}",
                    },
                }
            )

        image_data_uri = encode_jpeg_data_uri(image, quality=90)

        return {
            "image": image_data_uri,
            "rectangles": result_rects,
            "interpolation_mode": source.interpolation_mode,
        }
    except HTTPException:
        raise
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error("Key colors test failed: %s", e, exc_info=True)
        raise HTTPException(status_code=500, detail="Internal server error")
    finally:
        if stream:
            try:
                stream.stop()
            except Exception as e:
                logger.exception("unexpected in stream.stop cleanup: %s", e)


@router.websocket("/api/v1/color-strip-sources/{source_id}/key-colors/test/ws")
async def test_key_colors_ws(
    websocket: WebSocket,
    source_id: str,
    fps: int = Query(3),
    preview_width: int = Query(480),
):
    """WebSocket for real-time key_colors test preview with frame + rectangle overlay.

    Auth via first-message handshake: ``{"type":"auth","token":"..."}``
    within 3 s of connect.
    """
    import json as ws_json
    import time as ws_time
    from ledgrab.api.auth import accept_and_authenticate_ws
    from ledgrab.storage.color_strip_source import KeyColorsColorStripSource
    from ledgrab.core.capture.screen_capture import (
        calculate_average_color,
        calculate_dominant_color,
        calculate_median_color,
    )
    from ledgrab.storage.picture_source import ScreenCapturePictureSource
    from ledgrab.utils.image_codec import encode_jpeg_data_uri, resize_down

    if await accept_and_authenticate_ws(websocket) is None:
        return

    store = get_color_strip_store()
    source_store = get_picture_source_store()
    manager = get_processor_manager()
    device_store = get_device_store()

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

    if not isinstance(source, KeyColorsColorStripSource):
        await websocket.close(code=4003, reason="Not a key_colors source")
        return
    if not source.rectangles:
        await websocket.close(code=4003, reason="No regions configured")
        return
    if not source.picture_source_id:
        await websocket.close(code=4003, reason="No picture source configured")
        return

    try:
        chain = source_store.resolve_stream_chain(source.picture_source_id)
    except ValueError as e:
        await websocket.close(code=4003, reason=str(e))
        return

    raw_stream = chain["raw_stream"]
    if isinstance(raw_stream, ScreenCapturePictureSource):
        locked = manager.get_display_lock_info(raw_stream.display_index)
        if locked:
            try:
                name = device_store.get_device(locked).name
            except Exception:
                name = locked
            await websocket.close(code=4003, reason=f"Display captured by '{name}'")
            return

    fps = max(1, min(30, fps))
    preview_width = max(120, min(1920, preview_width))
    frame_interval = 1.0 / fps

    calc_fns = {
        "average": calculate_average_color,
        "median": calculate_median_color,
        "dominant": calculate_dominant_color,
    }
    calc_fn = calc_fns.get(source.interpolation_mode, calculate_average_color)

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

    # Use LiveStreamManager — reuses the already-running capture engine when a
    # target is active.  BetterCam (DXGI) only supports one session per display,
    # so creating a second engine would produce no frames.
    live_stream_mgr = manager._live_stream_manager
    try:
        live_stream = await asyncio.to_thread(live_stream_mgr.acquire, source.picture_source_id)
    except Exception as e:
        logger.error(f"KC test: LiveStream acquire failed: {e}")
        await websocket.send_text(ws_json.dumps({"type": "error", "detail": str(e)}))
        await websocket.close(code=4003, reason=str(e))
        return

    try:
        prev_frame_ref = None

        while True:
            loop_start = ws_time.monotonic()
            try:
                capture = await asyncio.to_thread(live_stream.get_latest_frame)
                if capture is None or capture.image is None:
                    await asyncio.sleep(frame_interval)
                    continue
                if capture is prev_frame_ref:
                    await asyncio.sleep(frame_interval * 0.5)
                    continue
                prev_frame_ref = capture
                cur_image = capture.image

                if not isinstance(cur_image, np.ndarray):
                    await asyncio.sleep(frame_interval)
                    continue

                # NOTE: postprocessing is already applied by the ProcessedLiveStream
                # when using LiveStreamManager.acquire() — do NOT re-apply here.

                # Re-read source for hot-update support
                try:
                    source = store.get_source(source_id)
                    calc_fn = calc_fns.get(source.interpolation_mode, calculate_average_color)
                except (KeyError, ValueError, AttributeError) as e:
                    logger.debug("source re-read failed for %s: %s", source_id, e)

                h, w = cur_image.shape[:2]
                result_rects = []
                for rect in source.rectangles:
                    px_x = max(0, int(rect.x * w))
                    px_y = max(0, int(rect.y * h))
                    px_w = max(1, int(rect.width * w))
                    px_h = max(1, int(rect.height * h))
                    px_x, px_y = min(px_x, w - 1), min(px_y, h - 1)
                    px_w, px_h = min(px_w, w - px_x), min(px_h, h - px_y)
                    sub = cur_image[px_y : px_y + px_h, px_x : px_x + px_w]
                    r, g, b = calc_fn(sub)
                    result_rects.append(
                        {
                            "name": rect.name,
                            "x": rect.x,
                            "y": rect.y,
                            "width": rect.width,
                            "height": rect.height,
                            "color": {
                                "r": int(r),
                                "g": int(g),
                                "b": int(b),
                                "hex": f"#{int(r):02x}{int(g):02x}{int(b):02x}",
                            },
                        }
                    )

                frame_to_encode = resize_down(cur_image, preview_width)
                frame_uri = encode_jpeg_data_uri(frame_to_encode, quality=85)

                await websocket.send_text(
                    ws_json.dumps(
                        {
                            "type": "frame",
                            "image": frame_uri,
                            "rectangles": result_rects,
                            "interpolation_mode": source.interpolation_mode,
                        }
                    )
                )

            except (WebSocketDisconnect, Exception) as inner_e:
                if isinstance(inner_e, WebSocketDisconnect):
                    raise
                logger.warning(f"KC CSS test WS frame error: {inner_e}")

            elapsed = ws_time.monotonic() - loop_start
            if frame_interval - elapsed > 0:
                await asyncio.sleep(frame_interval - elapsed)

    except WebSocketDisconnect:
        logger.info(f"KC CSS test WS disconnected for {source_id}")
    except Exception as e:
        logger.error(f"KC CSS test WS error: {e}", exc_info=True)
    finally:
        try:
            await asyncio.to_thread(live_stream_mgr.release, source.picture_source_id)
        except Exception as e:
            logger.exception("unexpected in live_stream_mgr.release: %s", e)
