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

"""Picture source routes."""

import base64
import io
import time

import httpx
import numpy as np
from PIL import Image
from fastapi import APIRouter, HTTPException, Depends, Query
from fastapi.responses import Response

from wled_controller.api.auth import AuthRequired
from wled_controller.api.dependencies import (
    get_device_store,
    get_picture_source_store,
    get_picture_target_store,
    get_pp_template_store,
    get_processor_manager,
    get_template_store,
)
from wled_controller.api.schemas.common import (
    CaptureImage,
    PerformanceMetrics,
    TemplateTestResponse,
)
from wled_controller.api.schemas.picture_sources import (
    ImageValidateRequest,
    ImageValidateResponse,
    PictureSourceCreate,
    PictureSourceListResponse,
    PictureSourceResponse,
    PictureSourceTestRequest,
    PictureSourceUpdate,
)
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.storage import DeviceStore
from wled_controller.storage.picture_target_store import PictureTargetStore
from wled_controller.storage.template_store import TemplateStore
from wled_controller.storage.postprocessing_template_store import PostprocessingTemplateStore
from wled_controller.storage.picture_source_store import PictureSourceStore
from wled_controller.storage.picture_source import ScreenCapturePictureSource, StaticImagePictureSource
from wled_controller.utils import get_logger

logger = get_logger(__name__)

router = APIRouter()


def _stream_to_response(s) -> PictureSourceResponse:
    """Convert a PictureSource to its API response."""
    return PictureSourceResponse(
        id=s.id,
        name=s.name,
        stream_type=s.stream_type,
        display_index=getattr(s, "display_index", None),
        capture_template_id=getattr(s, "capture_template_id", None),
        target_fps=getattr(s, "target_fps", None),
        source_stream_id=getattr(s, "source_stream_id", None),
        postprocessing_template_id=getattr(s, "postprocessing_template_id", None),
        image_source=getattr(s, "image_source", None),
        created_at=s.created_at,
        updated_at=s.updated_at,
        description=s.description,
    )


@router.get("/api/v1/picture-sources", response_model=PictureSourceListResponse, tags=["Picture Sources"])
async def list_picture_sources(
    _auth: AuthRequired,
    store: PictureSourceStore = Depends(get_picture_source_store),
):
    """List all picture sources."""
    try:
        streams = store.get_all_streams()
        responses = [_stream_to_response(s) for s in streams]
        return PictureSourceListResponse(streams=responses, count=len(responses))
    except Exception as e:
        logger.error(f"Failed to list picture sources: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/api/v1/picture-sources/validate-image", response_model=ImageValidateResponse, tags=["Picture Sources"])
async def validate_image(
    data: ImageValidateRequest,
    _auth: AuthRequired,
):
    """Validate an image source (URL or file path) and return a preview thumbnail."""
    try:
        from pathlib import Path

        source = data.image_source.strip()
        if not source:
            return ImageValidateResponse(valid=False, error="Image source is empty")

        if source.startswith(("http://", "https://")):
            async with httpx.AsyncClient(timeout=15, follow_redirects=True) as client:
                response = await client.get(source)
                response.raise_for_status()
                pil_image = Image.open(io.BytesIO(response.content))
        else:
            path = Path(source)
            if not path.exists():
                return ImageValidateResponse(valid=False, error=f"File not found: {source}")
            pil_image = Image.open(path)

        pil_image = pil_image.convert("RGB")
        width, height = pil_image.size

        # Create thumbnail preview (max 320px wide)
        thumb = pil_image.copy()
        thumb.thumbnail((320, 320), Image.Resampling.LANCZOS)
        buf = io.BytesIO()
        thumb.save(buf, format="JPEG", quality=80)
        buf.seek(0)
        preview = f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"

        return ImageValidateResponse(
            valid=True, width=width, height=height, preview=preview
        )

    except httpx.HTTPStatusError as e:
        return ImageValidateResponse(valid=False, error=f"HTTP {e.response.status_code}: {e.response.reason_phrase}")
    except httpx.RequestError as e:
        return ImageValidateResponse(valid=False, error=f"Request failed: {e}")
    except Exception as e:
        return ImageValidateResponse(valid=False, error=str(e))


@router.get("/api/v1/picture-sources/full-image", tags=["Picture Sources"])
async def get_full_image(
    _auth: AuthRequired,
    source: str = Query(..., description="Image URL or local file path"),
):
    """Serve the full-resolution image for lightbox preview."""
    from pathlib import Path

    try:
        if source.startswith(("http://", "https://")):
            async with httpx.AsyncClient(timeout=15, follow_redirects=True) as client:
                response = await client.get(source)
                response.raise_for_status()
                pil_image = Image.open(io.BytesIO(response.content))
        else:
            path = Path(source)
            if not path.exists():
                raise HTTPException(status_code=404, detail="File not found")
            pil_image = Image.open(path)

        pil_image = pil_image.convert("RGB")
        buf = io.BytesIO()
        pil_image.save(buf, format="JPEG", quality=90)
        buf.seek(0)
        return Response(content=buf.getvalue(), media_type="image/jpeg")

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.post("/api/v1/picture-sources", response_model=PictureSourceResponse, tags=["Picture Sources"], status_code=201)
async def create_picture_source(
    data: PictureSourceCreate,
    _auth: AuthRequired,
    store: PictureSourceStore = Depends(get_picture_source_store),
    template_store: TemplateStore = Depends(get_template_store),
    pp_store: PostprocessingTemplateStore = Depends(get_pp_template_store),
):
    """Create a new picture source."""
    try:
        # Validate referenced entities
        if data.stream_type == "raw" and data.capture_template_id:
            try:
                template_store.get_template(data.capture_template_id)
            except ValueError:
                raise HTTPException(
                    status_code=400,
                    detail=f"Capture template not found: {data.capture_template_id}",
                )

        if data.stream_type == "processed" and data.postprocessing_template_id:
            try:
                pp_store.get_template(data.postprocessing_template_id)
            except ValueError:
                raise HTTPException(
                    status_code=400,
                    detail=f"Postprocessing template not found: {data.postprocessing_template_id}",
                )

        stream = store.create_stream(
            name=data.name,
            stream_type=data.stream_type,
            display_index=data.display_index,
            capture_template_id=data.capture_template_id,
            target_fps=data.target_fps,
            source_stream_id=data.source_stream_id,
            postprocessing_template_id=data.postprocessing_template_id,
            image_source=data.image_source,
            description=data.description,
        )
        return _stream_to_response(stream)
    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 picture source: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/api/v1/picture-sources/{stream_id}", response_model=PictureSourceResponse, tags=["Picture Sources"])
async def get_picture_source(
    stream_id: str,
    _auth: AuthRequired,
    store: PictureSourceStore = Depends(get_picture_source_store),
):
    """Get picture source by ID."""
    try:
        stream = store.get_stream(stream_id)
        return _stream_to_response(stream)
    except ValueError:
        raise HTTPException(status_code=404, detail=f"Picture source {stream_id} not found")


@router.put("/api/v1/picture-sources/{stream_id}", response_model=PictureSourceResponse, tags=["Picture Sources"])
async def update_picture_source(
    stream_id: str,
    data: PictureSourceUpdate,
    _auth: AuthRequired,
    store: PictureSourceStore = Depends(get_picture_source_store),
):
    """Update a picture source."""
    try:
        stream = store.update_stream(
            stream_id=stream_id,
            name=data.name,
            display_index=data.display_index,
            capture_template_id=data.capture_template_id,
            target_fps=data.target_fps,
            source_stream_id=data.source_stream_id,
            postprocessing_template_id=data.postprocessing_template_id,
            image_source=data.image_source,
            description=data.description,
        )
        return _stream_to_response(stream)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to update picture source: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.delete("/api/v1/picture-sources/{stream_id}", status_code=204, tags=["Picture Sources"])
async def delete_picture_source(
    stream_id: str,
    _auth: AuthRequired,
    store: PictureSourceStore = Depends(get_picture_source_store),
    target_store: PictureTargetStore = Depends(get_picture_target_store),
):
    """Delete a picture source."""
    try:
        # Check if any target references this stream
        target_names = store.get_targets_referencing(stream_id, target_store)
        if target_names:
            names = ", ".join(target_names)
            raise HTTPException(
                status_code=409,
                detail=f"Cannot delete picture source: it is assigned to target(s): {names}. "
                       "Please reassign those targets before deleting.",
            )
        store.delete_stream(stream_id)
    except HTTPException:
        raise
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to delete picture source: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/api/v1/picture-sources/{stream_id}/test", response_model=TemplateTestResponse, tags=["Picture Sources"])
async def test_picture_source(
    stream_id: str,
    test_request: PictureSourceTestRequest,
    _auth: AuthRequired,
    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_store: PostprocessingTemplateStore = Depends(get_pp_template_store),
):
    """Test a picture source by resolving its chain and running a capture test.

    Resolves the stream chain to the raw stream, captures frames,
    and returns preview image + performance metrics.
    For processed streams, applies postprocessing (gamma, saturation, brightness)
    to the preview image.
    """
    stream = None
    try:
        # Resolve stream chain
        try:
            chain = store.resolve_stream_chain(stream_id)
        except ValueError as e:
            raise HTTPException(status_code=400, detail=str(e))

        raw_stream = chain["raw_stream"]

        if isinstance(raw_stream, StaticImagePictureSource):
            # Static image stream: load image directly, no engine needed
            from pathlib import Path

            source = raw_stream.image_source
            start_time = time.perf_counter()

            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:
                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")

            actual_duration = time.perf_counter() - start_time
            frame_count = 1
            total_capture_time = actual_duration

        elif isinstance(raw_stream, ScreenCapturePictureSource):
            # Screen capture stream: use engine
            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()

            frame_count = 0
            total_capture_time = 0.0
            last_frame = None
            start_time = time.perf_counter()

            if test_request.capture_duration == 0:
                # Single frame capture
                logger.info(f"Capturing single frame for {stream_id}")
                capture_start = time.perf_counter()
                screen_capture = stream.capture_frame()
                capture_elapsed = time.perf_counter() - capture_start
                if screen_capture is not None:
                    total_capture_time = capture_elapsed
                    frame_count = 1
                    last_frame = screen_capture
            else:
                logger.info(f"Starting {test_request.capture_duration}s stream test for {stream_id}")
                end_time = start_time + test_request.capture_duration
                while time.perf_counter() < end_time:
                    capture_start = time.perf_counter()
                    screen_capture = stream.capture_frame()
                    capture_elapsed = time.perf_counter() - capture_start
                    if screen_capture is None:
                        continue
                    total_capture_time += capture_elapsed
                    frame_count += 1
                    last_frame = screen_capture

            actual_duration = time.perf_counter() - start_time

            if last_frame is None:
                raise RuntimeError("No frames captured during test")

            if isinstance(last_frame.image, np.ndarray):
                pil_image = Image.fromarray(last_frame.image)
            else:
                raise ValueError("Unexpected image format from engine")

        # Create thumbnail
        thumbnail_width = 640
        aspect_ratio = pil_image.height / pil_image.width
        thumbnail_height = int(thumbnail_width * aspect_ratio)
        thumbnail = pil_image.copy()
        thumbnail.thumbnail((thumbnail_width, thumbnail_height), Image.Resampling.LANCZOS)

        # Apply postprocessing filters if this is a processed stream
        pp_template_ids = chain["postprocessing_template_ids"]
        if pp_template_ids:
            try:
                pp_template = pp_store.get_template(pp_template_ids[0])
                pool = ImagePool()

                def apply_filters(img):
                    arr = np.array(img)
                    for fi in pp_template.filters:
                        f = FilterRegistry.create_instance(fi.filter_id, fi.options)
                        result = f.process_image(arr, pool)
                        if result is not None:
                            arr = result
                    return Image.fromarray(arr)

                thumbnail = apply_filters(thumbnail)
                pil_image = apply_filters(pil_image)
            except ValueError:
                logger.warning(f"PP template {pp_template_ids[0]} not found, skipping postprocessing preview")

        # Encode thumbnail
        img_buffer = io.BytesIO()
        thumbnail.save(img_buffer, format='JPEG', quality=85)
        img_buffer.seek(0)
        thumbnail_b64 = base64.b64encode(img_buffer.getvalue()).decode('utf-8')
        thumbnail_data_uri = f"data:image/jpeg;base64,{thumbnail_b64}"

        # Encode full-resolution image
        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')
        full_data_uri = f"data:image/jpeg;base64,{full_b64}"

        actual_fps = frame_count / actual_duration if actual_duration > 0 else 0
        avg_capture_time_ms = (total_capture_time / frame_count * 1000) if frame_count > 0 else 0
        width, height = pil_image.size

        return TemplateTestResponse(
            full_capture=CaptureImage(
                image=thumbnail_data_uri,
                full_image=full_data_uri,
                width=width,
                height=height,
                thumbnail_width=thumbnail_width,
                thumbnail_height=thumbnail_height,
            ),
            border_extraction=None,
            performance=PerformanceMetrics(
                capture_duration_s=actual_duration,
                frame_count=frame_count,
                actual_fps=actual_fps,
                avg_capture_time_ms=avg_capture_time_ms,
            ),
        )

    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"Engine error: {str(e)}")
    except Exception as e:
        logger.error(f"Failed to test picture source: {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}")