wled-screen-controller-mixed/server/src/wled_controller/api/routes/output_targets.py

"""Output target routes: CRUD, processing control, settings, state, metrics."""

import asyncio
import base64
import io
import time

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

from wled_controller.api.auth import AuthRequired
from wled_controller.api.dependencies import (
    fire_entity_event,
    get_color_strip_store,
    get_device_store,
    get_pattern_template_store,
    get_picture_source_store,
    get_output_target_store,
    get_pp_template_store,
    get_processor_manager,
    get_template_store,
)
from wled_controller.api.schemas.output_targets import (
    BulkTargetRequest,
    BulkTargetResponse,
    ExtractedColorResponse,
    KCTestRectangleResponse,
    KCTestResponse,
    KeyColorsResponse,
    KeyColorsSettingsSchema,
    OutputTargetCreate,
    OutputTargetListResponse,
    OutputTargetResponse,
    OutputTargetUpdate,
    TargetMetricsResponse,
    TargetProcessingState,
)
from wled_controller.core.capture_engines import EngineRegistry
from wled_controller.core.filters import FilterRegistry, ImagePool
from wled_controller.core.processing.processor_manager import ProcessorManager
from wled_controller.core.capture.screen_capture import (
    calculate_average_color,
    calculate_dominant_color,
    calculate_median_color,
    get_available_displays,
)
from wled_controller.storage.color_strip_store import ColorStripStore
from wled_controller.storage.color_strip_source import AdvancedPictureColorStripSource, PictureColorStripSource
from wled_controller.storage import DeviceStore
from wled_controller.storage.pattern_template_store import PatternTemplateStore
from wled_controller.storage.picture_source import ScreenCapturePictureSource, StaticImagePictureSource
from wled_controller.storage.picture_source_store import PictureSourceStore
from wled_controller.storage.template_store import TemplateStore
from wled_controller.storage.wled_output_target import WledOutputTarget
from wled_controller.storage.key_colors_output_target import (
    KeyColorsSettings,
    KeyColorsOutputTarget,
)
from wled_controller.storage.output_target_store import OutputTargetStore
from wled_controller.utils import get_logger

logger = get_logger(__name__)

router = APIRouter()


def _kc_settings_to_schema(settings: KeyColorsSettings) -> KeyColorsSettingsSchema:
    """Convert core KeyColorsSettings to schema."""
    return KeyColorsSettingsSchema(
        fps=settings.fps,
        interpolation_mode=settings.interpolation_mode,
        smoothing=settings.smoothing,
        pattern_template_id=settings.pattern_template_id,
        brightness=settings.brightness,
        brightness_value_source_id=settings.brightness_value_source_id,
    )


def _kc_schema_to_settings(schema: KeyColorsSettingsSchema) -> KeyColorsSettings:
    """Convert schema KeyColorsSettings to core."""
    return KeyColorsSettings(
        fps=schema.fps,
        interpolation_mode=schema.interpolation_mode,
        smoothing=schema.smoothing,
        pattern_template_id=schema.pattern_template_id,
        brightness=schema.brightness,
        brightness_value_source_id=schema.brightness_value_source_id,
    )


def _target_to_response(target) -> OutputTargetResponse:
    """Convert an OutputTarget to OutputTargetResponse."""
    if isinstance(target, WledOutputTarget):
        return OutputTargetResponse(
            id=target.id,
            name=target.name,
            target_type=target.target_type,
            device_id=target.device_id,
            color_strip_source_id=target.color_strip_source_id,
            brightness_value_source_id=target.brightness_value_source_id or "",
            fps=target.fps,
            keepalive_interval=target.keepalive_interval,
            state_check_interval=target.state_check_interval,
            min_brightness_threshold=target.min_brightness_threshold,
            adaptive_fps=target.adaptive_fps,
            protocol=target.protocol,
            description=target.description,
            tags=getattr(target, 'tags', []),

            created_at=target.created_at,
            updated_at=target.updated_at,
        )
    elif isinstance(target, KeyColorsOutputTarget):
        return OutputTargetResponse(
            id=target.id,
            name=target.name,
            target_type=target.target_type,
            picture_source_id=target.picture_source_id,
            key_colors_settings=_kc_settings_to_schema(target.settings),
            description=target.description,
            tags=getattr(target, 'tags', []),

            created_at=target.created_at,
            updated_at=target.updated_at,
        )
    else:
        return OutputTargetResponse(
            id=target.id,
            name=target.name,
            target_type=target.target_type,
            description=target.description,
            tags=getattr(target, 'tags', []),

            created_at=target.created_at,
            updated_at=target.updated_at,
        )


# ===== CRUD ENDPOINTS =====

@router.post("/api/v1/output-targets", response_model=OutputTargetResponse, tags=["Targets"], status_code=201)
async def create_target(
    data: OutputTargetCreate,
    _auth: AuthRequired,
    target_store: OutputTargetStore = Depends(get_output_target_store),
    device_store: DeviceStore = Depends(get_device_store),
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Create a new output target."""
    try:
        # Validate device exists if provided
        if data.device_id:
            try:
                device_store.get_device(data.device_id)
            except ValueError:
                raise HTTPException(status_code=422, detail=f"Device {data.device_id} not found")

        kc_settings = _kc_schema_to_settings(data.key_colors_settings) if data.key_colors_settings else None

        # Create in store
        target = target_store.create_target(
            name=data.name,
            target_type=data.target_type,
            device_id=data.device_id,
            color_strip_source_id=data.color_strip_source_id,
            brightness_value_source_id=data.brightness_value_source_id,
            fps=data.fps,
            keepalive_interval=data.keepalive_interval,
            state_check_interval=data.state_check_interval,
            min_brightness_threshold=data.min_brightness_threshold,
            adaptive_fps=data.adaptive_fps,
            protocol=data.protocol,
            picture_source_id=data.picture_source_id,
            key_colors_settings=kc_settings,
            description=data.description,
            tags=data.tags,
        )

        # Register in processor manager
        try:
            target.register_with_manager(manager)
        except ValueError as e:
            logger.warning(f"Could not register target {target.id} in processor manager: {e}")

        fire_entity_event("output_target", "created", target.id)
        return _target_to_response(target)

    except HTTPException:
        raise
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to create target: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/api/v1/output-targets", response_model=OutputTargetListResponse, tags=["Targets"])
async def list_targets(
    _auth: AuthRequired,
    target_store: OutputTargetStore = Depends(get_output_target_store),
):
    """List all output targets."""
    targets = target_store.get_all_targets()
    responses = [_target_to_response(t) for t in targets]
    return OutputTargetListResponse(targets=responses, count=len(responses))


@router.get("/api/v1/output-targets/batch/states", tags=["Processing"])
async def batch_target_states(
    _auth: AuthRequired,
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Get processing state for all targets in a single request."""
    return {"states": manager.get_all_target_states()}


@router.get("/api/v1/output-targets/batch/metrics", tags=["Metrics"])
async def batch_target_metrics(
    _auth: AuthRequired,
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Get metrics for all targets in a single request."""
    return {"metrics": manager.get_all_target_metrics()}


@router.get("/api/v1/output-targets/{target_id}", response_model=OutputTargetResponse, tags=["Targets"])
async def get_target(
    target_id: str,
    _auth: AuthRequired,
    target_store: OutputTargetStore = Depends(get_output_target_store),
):
    """Get a output target by ID."""
    try:
        target = target_store.get_target(target_id)
        return _target_to_response(target)
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))


@router.put("/api/v1/output-targets/{target_id}", response_model=OutputTargetResponse, tags=["Targets"])
async def update_target(
    target_id: str,
    data: OutputTargetUpdate,
    _auth: AuthRequired,
    target_store: OutputTargetStore = Depends(get_output_target_store),
    device_store: DeviceStore = Depends(get_device_store),
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Update a output target."""
    try:
        # Validate device exists if changing
        if data.device_id is not None and data.device_id:
            try:
                device_store.get_device(data.device_id)
            except ValueError:
                raise HTTPException(status_code=422, detail=f"Device {data.device_id} not found")

        # Build KC settings with partial-update support: only apply fields that were
        # explicitly provided in the request body, merging with the existing settings.
        kc_settings = None
        if data.key_colors_settings is not None:
            incoming = data.key_colors_settings.model_dump(exclude_unset=True)
            try:
                existing_target = target_store.get_target(target_id)
            except ValueError:
                existing_target = None

            if isinstance(existing_target, KeyColorsOutputTarget):
                ex = existing_target.settings
                merged = KeyColorsSettingsSchema(
                    fps=incoming.get("fps", ex.fps),
                    interpolation_mode=incoming.get("interpolation_mode", ex.interpolation_mode),
                    smoothing=incoming.get("smoothing", ex.smoothing),
                    pattern_template_id=incoming.get("pattern_template_id", ex.pattern_template_id),
                    brightness=incoming.get("brightness", ex.brightness),
                    brightness_value_source_id=incoming.get("brightness_value_source_id", ex.brightness_value_source_id),
                )
                kc_settings = _kc_schema_to_settings(merged)
            else:
                kc_settings = _kc_schema_to_settings(data.key_colors_settings)

        # Update in store
        target = target_store.update_target(
            target_id=target_id,
            name=data.name,
            device_id=data.device_id,
            color_strip_source_id=data.color_strip_source_id,
            brightness_value_source_id=data.brightness_value_source_id,
            fps=data.fps,
            keepalive_interval=data.keepalive_interval,
            state_check_interval=data.state_check_interval,
            min_brightness_threshold=data.min_brightness_threshold,
            adaptive_fps=data.adaptive_fps,
            protocol=data.protocol,
            key_colors_settings=kc_settings,
            description=data.description,
            tags=data.tags,
        )

        # Detect KC brightness VS change (inside key_colors_settings)
        kc_brightness_vs_changed = False
        if data.key_colors_settings is not None:
            kc_incoming = data.key_colors_settings.model_dump(exclude_unset=True)
            if "brightness_value_source_id" in kc_incoming:
                kc_brightness_vs_changed = True

        # Sync processor manager (run in thread — css release/acquire can block)
        try:
            await asyncio.to_thread(
                target.sync_with_manager,
                manager,
                settings_changed=(data.fps is not None or
                                  data.keepalive_interval is not None or
                                  data.state_check_interval is not None or
                                  data.min_brightness_threshold is not None or
                                  data.adaptive_fps is not None or
                                  data.key_colors_settings is not None),
                css_changed=data.color_strip_source_id is not None,
                brightness_vs_changed=(data.brightness_value_source_id is not None or kc_brightness_vs_changed),
            )
        except ValueError:
            pass

        # Device change requires async stop → swap → start cycle
        if data.device_id is not None:
            try:
                await manager.update_target_device(target_id, target.device_id)
            except ValueError:
                pass

        fire_entity_event("output_target", "updated", target_id)
        return _target_to_response(target)

    except HTTPException:
        raise
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to update target: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.delete("/api/v1/output-targets/{target_id}", status_code=204, tags=["Targets"])
async def delete_target(
    target_id: str,
    _auth: AuthRequired,
    target_store: OutputTargetStore = Depends(get_output_target_store),
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Delete a output target. Stops processing first if active."""
    try:
        # Stop processing if running
        try:
            await manager.stop_processing(target_id)
        except ValueError:
            pass

        # Remove from manager
        try:
            manager.remove_target(target_id)
        except (ValueError, RuntimeError):
            pass

        # Delete from store
        target_store.delete_target(target_id)

        fire_entity_event("output_target", "deleted", target_id)
        logger.info(f"Deleted target {target_id}")

    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to delete target: {e}")
        raise HTTPException(status_code=500, detail=str(e))


# ===== BULK PROCESSING CONTROL ENDPOINTS =====

@router.post("/api/v1/output-targets/bulk/start", response_model=BulkTargetResponse, tags=["Processing"])
async def bulk_start_processing(
    body: BulkTargetRequest,
    _auth: AuthRequired,
    target_store: OutputTargetStore = Depends(get_output_target_store),
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Start processing for multiple output targets. Returns lists of started IDs and per-ID errors."""
    started: list[str] = []
    errors: dict[str, str] = {}

    for target_id in body.ids:
        try:
            target_store.get_target(target_id)
            await manager.start_processing(target_id)
            started.append(target_id)
            logger.info(f"Bulk start: started processing for target {target_id}")
        except ValueError as e:
            errors[target_id] = str(e)
        except RuntimeError as e:
            msg = str(e)
            for t in target_store.get_all_targets():
                if t.id in msg:
                    msg = msg.replace(t.id, f"'{t.name}'")
            errors[target_id] = msg
        except Exception as e:
            logger.error(f"Bulk start: failed to start target {target_id}: {e}")
            errors[target_id] = str(e)

    return BulkTargetResponse(started=started, errors=errors)


@router.post("/api/v1/output-targets/bulk/stop", response_model=BulkTargetResponse, tags=["Processing"])
async def bulk_stop_processing(
    body: BulkTargetRequest,
    _auth: AuthRequired,
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Stop processing for multiple output targets. Returns lists of stopped IDs and per-ID errors."""
    stopped: list[str] = []
    errors: dict[str, str] = {}

    for target_id in body.ids:
        try:
            await manager.stop_processing(target_id)
            stopped.append(target_id)
            logger.info(f"Bulk stop: stopped processing for target {target_id}")
        except ValueError as e:
            errors[target_id] = str(e)
        except Exception as e:
            logger.error(f"Bulk stop: failed to stop target {target_id}: {e}")
            errors[target_id] = str(e)

    return BulkTargetResponse(stopped=stopped, errors=errors)


# ===== PROCESSING CONTROL ENDPOINTS =====

@router.post("/api/v1/output-targets/{target_id}/start", tags=["Processing"])
async def start_processing(
    target_id: str,
    _auth: AuthRequired,
    target_store: OutputTargetStore = Depends(get_output_target_store),
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Start processing for a output target."""
    try:
        # Verify target exists in store
        target_store.get_target(target_id)

        await manager.start_processing(target_id)

        logger.info(f"Started processing for target {target_id}")
        return {"status": "started", "target_id": target_id}

    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except RuntimeError as e:
        # Resolve target IDs to human-readable names in error messages
        msg = str(e)
        for t in target_store.get_all_targets():
            if t.id in msg:
                msg = msg.replace(t.id, f"'{t.name}'")
        raise HTTPException(status_code=409, detail=msg)
    except Exception as e:
        logger.error(f"Failed to start processing: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/api/v1/output-targets/{target_id}/stop", tags=["Processing"])
async def stop_processing(
    target_id: str,
    _auth: AuthRequired,
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Stop processing for a output target."""
    try:
        await manager.stop_processing(target_id)

        logger.info(f"Stopped processing for target {target_id}")
        return {"status": "stopped", "target_id": target_id}

    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to stop processing: {e}")
        raise HTTPException(status_code=500, detail=str(e))


# ===== STATE & METRICS ENDPOINTS =====

@router.get("/api/v1/output-targets/{target_id}/state", response_model=TargetProcessingState, tags=["Processing"])
async def get_target_state(
    target_id: str,
    _auth: AuthRequired,
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Get current processing state for a target."""
    try:
        state = manager.get_target_state(target_id)
        return TargetProcessingState(**state)

    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to get target state: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/api/v1/output-targets/{target_id}/metrics", response_model=TargetMetricsResponse, tags=["Metrics"])
async def get_target_metrics(
    target_id: str,
    _auth: AuthRequired,
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Get processing metrics for a target."""
    try:
        metrics = manager.get_target_metrics(target_id)
        return TargetMetricsResponse(**metrics)

    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to get target metrics: {e}")
        raise HTTPException(status_code=500, detail=str(e))


# ===== KEY COLORS ENDPOINTS =====

@router.get("/api/v1/output-targets/{target_id}/colors", response_model=KeyColorsResponse, tags=["Key Colors"])
async def get_target_colors(
    target_id: str,
    _auth: AuthRequired,
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Get latest extracted colors for a key-colors target (polling)."""
    try:
        raw_colors = manager.get_kc_latest_colors(target_id)
        colors = {}
        for name, (r, g, b) in raw_colors.items():
            colors[name] = ExtractedColorResponse(
                r=r, g=g, b=b,
                hex=f"#{r:02x}{g:02x}{b:02x}",
            )
        from datetime import datetime, timezone
        return KeyColorsResponse(
            target_id=target_id,
            colors=colors,
            timestamp=datetime.now(timezone.utc),
        )
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))


@router.post("/api/v1/output-targets/{target_id}/test", response_model=KCTestResponse, tags=["Key Colors"])
async def test_kc_target(
    target_id: str,
    _auth: AuthRequired,
    target_store: OutputTargetStore = Depends(get_output_target_store),
    source_store: PictureSourceStore = Depends(get_picture_source_store),
    template_store: TemplateStore = Depends(get_template_store),
    pattern_store: PatternTemplateStore = Depends(get_pattern_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 target: capture a frame, extract colors from each rectangle."""
    import httpx

    stream = None
    try:
        # 1. Load and validate KC target
        try:
            target = target_store.get_target(target_id)
        except ValueError as e:
            raise HTTPException(status_code=404, detail=str(e))

        if not isinstance(target, KeyColorsOutputTarget):
            raise HTTPException(status_code=400, detail="Target is not a key_colors target")

        settings = target.settings

        # 2. Resolve pattern template
        if not settings.pattern_template_id:
            raise HTTPException(status_code=400, detail="No pattern template configured")

        try:
            pattern_tmpl = pattern_store.get_template(settings.pattern_template_id)
        except ValueError:
            raise HTTPException(status_code=400, detail=f"Pattern template not found: {settings.pattern_template_id}")

        rectangles = pattern_tmpl.rectangles
        if not rectangles:
            raise HTTPException(status_code=400, detail="Pattern template has no rectangles")

        # 3. Resolve picture source and capture a frame
        if not target.picture_source_id:
            raise HTTPException(status_code=400, detail="No picture source configured")

        try:
            chain = source_store.resolve_stream_chain(target.picture_source_id)
        except ValueError as e:
            raise HTTPException(status_code=400, detail=str(e))

        raw_stream = chain["raw_stream"]

        if isinstance(raw_stream, StaticImagePictureSource):
            source = raw_stream.image_source
            if source.startswith(("http://", "https://")):
                async with httpx.AsyncClient(timeout=15, follow_redirects=True) as client:
                    resp = await client.get(source)
                    resp.raise_for_status()
                    pil_image = Image.open(io.BytesIO(resp.content)).convert("RGB")
            else:
                from pathlib import Path
                path = Path(source)
                if not path.exists():
                    raise HTTPException(status_code=400, detail=f"Image file not found: {source}")
                pil_image = Image.open(path).convert("RGB")

        elif isinstance(raw_stream, ScreenCapturePictureSource):
            try:
                capture_template = template_store.get_template(raw_stream.capture_template_id)
            except ValueError:
                raise HTTPException(
                    status_code=400,
                    detail=f"Capture template not found: {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}' is not available on this system",
                )

            locked_device_id = processor_manager.get_display_lock_info(display_index)
            if locked_device_id:
                try:
                    device = device_store.get_device(locked_device_id)
                    device_name = device.name
                except Exception:
                    device_name = locked_device_id
                raise HTTPException(
                    status_code=409,
                    detail=f"Display {display_index} is currently being captured by device '{device_name}'. "
                           f"Please stop the device processing before testing.",
                )

            stream = EngineRegistry.create_stream(
                capture_template.engine_type, display_index, capture_template.engine_config
            )
            stream.initialize()

            screen_capture = stream.capture_frame()
            if screen_capture is None:
                raise RuntimeError("No frame captured")

            if isinstance(screen_capture.image, np.ndarray):
                pil_image = Image.fromarray(screen_capture.image)
            else:
                raise ValueError("Unexpected image format from engine")
        else:
            raise HTTPException(status_code=400, detail="Unsupported picture source type")

        # 3b. Apply postprocessing filters (if the picture source has a filter chain)
        pp_template_ids = chain.get("postprocessing_template_ids", [])
        if pp_template_ids and pp_template_store:
            img_array = np.array(pil_image)
            image_pool = ImagePool()
            for pp_id in pp_template_ids:
                try:
                    pp_template = pp_template_store.get_template(pp_id)
                except ValueError:
                    logger.warning(f"KC test: PP template {pp_id} not found, skipping")
                    continue
                flat_filters = pp_template_store.resolve_filter_instances(pp_template.filters)
                for fi in flat_filters:
                    try:
                        f = FilterRegistry.create_instance(fi.filter_id, fi.options)
                        result = f.process_image(img_array, image_pool)
                        if result is not None:
                            img_array = result
                    except ValueError:
                        logger.warning(f"KC test: unknown filter '{fi.filter_id}', skipping")
            pil_image = Image.fromarray(img_array)

        # 4. Extract colors from each rectangle
        img_array = np.array(pil_image)
        h, w = img_array.shape[:2]

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

        result_rects = []
        for rect in 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 = min(px_x, w - 1)
            px_y = min(px_y, h - 1)
            px_w = min(px_w, w - px_x)
            px_h = min(px_h, h - px_y)

            sub_img = img_array[px_y:px_y + px_h, px_x:px_x + px_w]
            r, g, b = calc_fn(sub_img)

            result_rects.append(KCTestRectangleResponse(
                name=rect.name,
                x=rect.x,
                y=rect.y,
                width=rect.width,
                height=rect.height,
                color=ExtractedColorResponse(r=r, g=g, b=b, hex=f"#{r:02x}{g:02x}{b:02x}"),
            ))

        # 5. Encode frame as base64 JPEG
        full_buffer = io.BytesIO()
        pil_image.save(full_buffer, format='JPEG', quality=90)
        full_buffer.seek(0)
        full_b64 = base64.b64encode(full_buffer.getvalue()).decode('utf-8')
        image_data_uri = f"data:image/jpeg;base64,{full_b64}"

        return KCTestResponse(
            image=image_data_uri,
            rectangles=result_rects,
            interpolation_mode=settings.interpolation_mode,
            pattern_template_name=pattern_tmpl.name,
        )

    except HTTPException:
        raise
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except RuntimeError as e:
        raise HTTPException(status_code=500, detail=f"Capture error: {str(e)}")
    except Exception as e:
        logger.error(f"Failed to test KC target: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if stream:
            try:
                stream.cleanup()
            except Exception as e:
                logger.error(f"Error cleaning up test stream: {e}")


@router.websocket("/api/v1/output-targets/{target_id}/test/ws")
async def test_kc_target_ws(
    websocket: WebSocket,
    target_id: str,
    token: str = Query(""),
    fps: int = Query(3),
    preview_width: int = Query(480),
):
    """WebSocket for real-time KC target test preview. Auth via ?token=<api_key>.

    Streams JSON frames: {"type": "frame", "image": "data:image/jpeg;base64,...",
    "rectangles": [...], "pattern_template_name": "...", "interpolation_mode": "..."}
    """
    import json as _json
    from wled_controller.api.auth import verify_ws_token

    if not verify_ws_token(token):
        await websocket.close(code=4001, reason="Unauthorized")
        return

    # Load stores
    target_store_inst: OutputTargetStore = get_output_target_store()
    source_store_inst: PictureSourceStore = get_picture_source_store()
    template_store_inst: TemplateStore = get_template_store()
    pattern_store_inst: PatternTemplateStore = get_pattern_template_store()
    processor_manager_inst: ProcessorManager = get_processor_manager()
    device_store_inst: DeviceStore = get_device_store()
    pp_template_store_inst = get_pp_template_store()

    # Validate target
    try:
        target = target_store_inst.get_target(target_id)
    except ValueError as e:
        await websocket.close(code=4004, reason=str(e))
        return

    if not isinstance(target, KeyColorsOutputTarget):
        await websocket.close(code=4003, reason="Target is not a key_colors target")
        return

    settings = target.settings

    if not settings.pattern_template_id:
        await websocket.close(code=4003, reason="No pattern template configured")
        return

    try:
        pattern_tmpl = pattern_store_inst.get_template(settings.pattern_template_id)
    except ValueError:
        await websocket.close(code=4003, reason=f"Pattern template not found: {settings.pattern_template_id}")
        return

    rectangles = pattern_tmpl.rectangles
    if not rectangles:
        await websocket.close(code=4003, reason="Pattern template has no rectangles")
        return

    if not target.picture_source_id:
        await websocket.close(code=4003, reason="No picture source configured")
        return

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

    raw_stream = chain["raw_stream"]

    # For screen capture sources, check display lock
    if isinstance(raw_stream, ScreenCapturePictureSource):
        display_index = raw_stream.display_index
        locked_device_id = processor_manager_inst.get_display_lock_info(display_index)
        if locked_device_id:
            try:
                device = device_store_inst.get_device(locked_device_id)
                device_name = device.name
            except Exception:
                device_name = locked_device_id
            await websocket.close(
                code=4003,
                reason=f"Display {display_index} is captured by '{device_name}'. Stop processing first.",
            )
            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(settings.interpolation_mode, calculate_average_color)

    await websocket.accept()
    logger.info(f"KC test WS connected for {target_id} (fps={fps})")

    capture_stream = None
    try:
        while True:
            loop_start = time.monotonic()

            pil_image = None
            capture_stream_local = None

            try:
                import httpx
                # Reload chain each iteration for dynamic sources
                chain = source_store_inst.resolve_stream_chain(target.picture_source_id)
                raw_stream = chain["raw_stream"]

                if isinstance(raw_stream, StaticImagePictureSource):
                    source = raw_stream.image_source
                    if source.startswith(("http://", "https://")):
                        async with httpx.AsyncClient(timeout=15, follow_redirects=True) as client:
                            resp = await client.get(source)
                            resp.raise_for_status()
                            pil_image = Image.open(io.BytesIO(resp.content)).convert("RGB")
                    else:
                        from pathlib import Path
                        path = Path(source)
                        if path.exists():
                            pil_image = Image.open(path).convert("RGB")

                elif isinstance(raw_stream, ScreenCapturePictureSource):
                    try:
                        capture_tmpl = template_store_inst.get_template(raw_stream.capture_template_id)
                    except ValueError:
                        break

                    if capture_tmpl.engine_type not in EngineRegistry.get_available_engines():
                        break

                    capture_stream_local = EngineRegistry.create_stream(
                        capture_tmpl.engine_type, raw_stream.display_index, capture_tmpl.engine_config
                    )
                    capture_stream_local.initialize()
                    screen_capture = capture_stream_local.capture_frame()
                    if screen_capture is not None and isinstance(screen_capture.image, np.ndarray):
                        pil_image = Image.fromarray(screen_capture.image)

                else:
                    # VideoCaptureSource or other — not directly supported in WS test
                    break

                if pil_image is None:
                    await asyncio.sleep(frame_interval)
                    continue

                # Apply postprocessing
                pp_template_ids = chain.get("postprocessing_template_ids", [])
                if pp_template_ids and pp_template_store_inst:
                    img_array = np.array(pil_image)
                    image_pool = ImagePool()
                    for pp_id in pp_template_ids:
                        try:
                            pp_template = pp_template_store_inst.get_template(pp_id)
                        except ValueError:
                            continue
                        flat_filters = pp_template_store_inst.resolve_filter_instances(pp_template.filters)
                        for fi in flat_filters:
                            try:
                                f = FilterRegistry.create_instance(fi.filter_id, fi.options)
                                result = f.process_image(img_array, image_pool)
                                if result is not None:
                                    img_array = result
                            except ValueError:
                                pass
                    pil_image = Image.fromarray(img_array)

                # Extract colors
                img_array = np.array(pil_image)
                h, w = img_array.shape[:2]

                result_rects = []
                for rect in 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 = min(px_x, w - 1)
                    px_y = min(px_y, h - 1)
                    px_w = min(px_w, w - px_x)
                    px_h = min(px_h, h - px_y)

                    sub_img = img_array[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": r, "g": g, "b": b, "hex": f"#{r:02x}{g:02x}{b:02x}"},
                    })

                # Encode frame as JPEG
                if preview_width and pil_image.width > preview_width:
                    ratio = preview_width / pil_image.width
                    thumb = pil_image.resize((preview_width, int(pil_image.height * ratio)), Image.LANCZOS)
                else:
                    thumb = pil_image
                buf = io.BytesIO()
                thumb.save(buf, format="JPEG", quality=85)
                b64 = base64.b64encode(buf.getvalue()).decode()

                await websocket.send_text(_json.dumps({
                    "type": "frame",
                    "image": f"data:image/jpeg;base64,{b64}",
                    "rectangles": result_rects,
                    "pattern_template_name": pattern_tmpl.name,
                    "interpolation_mode": settings.interpolation_mode,
                }))

            except (WebSocketDisconnect, Exception) as inner_e:
                if isinstance(inner_e, WebSocketDisconnect):
                    raise
                logger.warning(f"KC test WS frame error for {target_id}: {inner_e}")
            finally:
                if capture_stream_local:
                    try:
                        capture_stream_local.cleanup()
                    except Exception:
                        pass

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

    except WebSocketDisconnect:
        logger.info(f"KC test WS disconnected for {target_id}")
    except Exception as e:
        logger.error(f"KC test WS error for {target_id}: {e}", exc_info=True)
    finally:
        if capture_stream:
            try:
                capture_stream.cleanup()
            except Exception:
                pass
        logger.info(f"KC test WS closed for {target_id}")


@router.websocket("/api/v1/output-targets/{target_id}/ws")
async def target_colors_ws(
    websocket: WebSocket,
    target_id: str,
    token: str = Query(""),
):
    """WebSocket for real-time key color updates. Auth via ?token=<api_key>."""
    from wled_controller.api.auth import verify_ws_token
    if not verify_ws_token(token):
        await websocket.close(code=4001, reason="Unauthorized")
        return

    await websocket.accept()

    manager = get_processor_manager()

    try:
        manager.add_kc_ws_client(target_id, websocket)
    except ValueError:
        await websocket.close(code=4004, reason="Target not found")
        return

    try:
        while True:
            # Keep alive — wait for client messages (or disconnect)
            await websocket.receive_text()
    except WebSocketDisconnect:
        pass
    finally:
        manager.remove_kc_ws_client(target_id, websocket)


@router.websocket("/api/v1/output-targets/{target_id}/led-preview/ws")
async def led_preview_ws(
    websocket: WebSocket,
    target_id: str,
    token: str = Query(""),
):
    """WebSocket for real-time LED strip preview. Sends binary RGB frames. Auth via ?token=<api_key>."""
    from wled_controller.api.auth import verify_ws_token
    if not verify_ws_token(token):
        await websocket.close(code=4001, reason="Unauthorized")
        return

    await websocket.accept()

    manager = get_processor_manager()

    try:
        manager.add_led_preview_client(target_id, websocket)
    except ValueError:
        await websocket.close(code=4004, reason="Target not found")
        return

    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        pass
    finally:
        manager.remove_led_preview_client(target_id, websocket)


# ===== STATE CHANGE EVENT STREAM =====


@router.websocket("/api/v1/events/ws")
async def events_ws(
    websocket: WebSocket,
    token: str = Query(""),
):
    """WebSocket for real-time state change events. Auth via ?token=<api_key>."""
    from wled_controller.api.auth import verify_ws_token
    if not verify_ws_token(token):
        await websocket.close(code=4001, reason="Unauthorized")
        return

    await websocket.accept()

    manager = get_processor_manager()
    queue = manager.subscribe_events()

    try:
        while True:
            event = await queue.get()
            await websocket.send_json(event)
    except WebSocketDisconnect:
        pass
    except Exception:
        pass
    finally:
        manager.unsubscribe_events(queue)


# ===== OVERLAY VISUALIZATION =====

@router.post("/api/v1/output-targets/{target_id}/overlay/start", tags=["Visualization"])
async def start_target_overlay(
    target_id: str,
    _auth: AuthRequired,
    manager: ProcessorManager = Depends(get_processor_manager),
    target_store: OutputTargetStore = Depends(get_output_target_store),
    color_strip_store: ColorStripStore = Depends(get_color_strip_store),
    picture_source_store: PictureSourceStore = Depends(get_picture_source_store),
):
    """Start screen overlay visualization for a target.

    Displays a transparent overlay on the target display showing:
    - Border sampling zones (colored rectangles)
    - LED position markers (numbered dots)
    - Pixel-to-LED mapping ranges (colored segments)
    - Calibration info text
    """
    try:
        # Get target name from store
        target = target_store.get_target(target_id)
        if not target:
            raise ValueError(f"Target {target_id} not found")

        # Pre-load calibration and display info from the CSS store so the overlay
        # can start even when processing is not currently running.
        calibration = None
        display_info = None
        if isinstance(target, WledOutputTarget) and target.color_strip_source_id:
            first_css_id = target.color_strip_source_id
            if first_css_id:
                try:
                    css = color_strip_store.get_source(first_css_id)
                    if isinstance(css, (PictureColorStripSource, AdvancedPictureColorStripSource)) and css.calibration:
                        calibration = css.calibration
                        # Resolve the display this CSS is capturing
                        from wled_controller.api.routes.color_strip_sources import _resolve_display_index
                        ps_id = getattr(css, "picture_source_id", "") or ""
                        display_index = _resolve_display_index(ps_id, picture_source_store)
                        displays = get_available_displays()
                        if displays:
                            display_index = min(display_index, len(displays) - 1)
                            display_info = displays[display_index]
                except Exception as e:
                    logger.warning(f"Could not pre-load CSS calibration for overlay on {target_id}: {e}")

        await manager.start_overlay(target_id, target.name, calibration=calibration, display_info=display_info)
        return {"status": "started", "target_id": target_id}

    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except RuntimeError as e:
        raise HTTPException(status_code=409, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to start overlay: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/api/v1/output-targets/{target_id}/overlay/stop", tags=["Visualization"])
async def stop_target_overlay(
    target_id: str,
    _auth: AuthRequired,
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Stop screen overlay visualization for a target."""
    try:
        await manager.stop_overlay(target_id)
        return {"status": "stopped", "target_id": target_id}

    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to stop overlay: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/api/v1/output-targets/{target_id}/overlay/status", tags=["Visualization"])
async def get_overlay_status(
    target_id: str,
    _auth: AuthRequired,
    manager: ProcessorManager = Depends(get_processor_manager),
):
    """Check if overlay is active for a target."""
    try:
        active = manager.is_overlay_active(target_id)
        return {"target_id": target_id, "active": active}

    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))