ledgrab/server/src/wled_controller/api/routes/postprocessing.py

"""Postprocessing template routes."""

import base64
import io
import time

import httpx
import numpy as np
from PIL import Image
from fastapi import APIRouter, HTTPException, Depends

from wled_controller.api.auth import AuthRequired
from wled_controller.api.dependencies import (
    get_device_store,
    get_picture_source_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.filters import FilterInstanceSchema
from wled_controller.api.schemas.postprocessing import (
    PostprocessingTemplateCreate,
    PostprocessingTemplateListResponse,
    PostprocessingTemplateResponse,
    PostprocessingTemplateUpdate,
    PPTemplateTestRequest,
)
from wled_controller.core.capture_engines import EngineRegistry
from wled_controller.core.filters import FilterRegistry, FilterInstance, ImagePool
from wled_controller.core.processor_manager import ProcessorManager
from wled_controller.storage import DeviceStore
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 _pp_template_to_response(t) -> PostprocessingTemplateResponse:
    """Convert a PostprocessingTemplate to its API response."""
    return PostprocessingTemplateResponse(
        id=t.id,
        name=t.name,
        filters=[FilterInstanceSchema(filter_id=f.filter_id, options=f.options) for f in t.filters],
        created_at=t.created_at,
        updated_at=t.updated_at,
        description=t.description,
    )


@router.get("/api/v1/postprocessing-templates", response_model=PostprocessingTemplateListResponse, tags=["Postprocessing Templates"])
async def list_pp_templates(
    _auth: AuthRequired,
    store: PostprocessingTemplateStore = Depends(get_pp_template_store),
):
    """List all postprocessing templates."""
    try:
        templates = store.get_all_templates()
        responses = [_pp_template_to_response(t) for t in templates]
        return PostprocessingTemplateListResponse(templates=responses, count=len(responses))
    except Exception as e:
        logger.error(f"Failed to list postprocessing templates: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/api/v1/postprocessing-templates", response_model=PostprocessingTemplateResponse, tags=["Postprocessing Templates"], status_code=201)
async def create_pp_template(
    data: PostprocessingTemplateCreate,
    _auth: AuthRequired,
    store: PostprocessingTemplateStore = Depends(get_pp_template_store),
):
    """Create a new postprocessing template."""
    try:
        filters = [FilterInstance(f.filter_id, f.options) for f in data.filters]
        template = store.create_template(
            name=data.name,
            filters=filters,
            description=data.description,
        )
        return _pp_template_to_response(template)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to create postprocessing template: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/api/v1/postprocessing-templates/{template_id}", response_model=PostprocessingTemplateResponse, tags=["Postprocessing Templates"])
async def get_pp_template(
    template_id: str,
    _auth: AuthRequired,
    store: PostprocessingTemplateStore = Depends(get_pp_template_store),
):
    """Get postprocessing template by ID."""
    try:
        template = store.get_template(template_id)
        return _pp_template_to_response(template)
    except ValueError:
        raise HTTPException(status_code=404, detail=f"Postprocessing template {template_id} not found")


@router.put("/api/v1/postprocessing-templates/{template_id}", response_model=PostprocessingTemplateResponse, tags=["Postprocessing Templates"])
async def update_pp_template(
    template_id: str,
    data: PostprocessingTemplateUpdate,
    _auth: AuthRequired,
    store: PostprocessingTemplateStore = Depends(get_pp_template_store),
):
    """Update a postprocessing template."""
    try:
        filters = [FilterInstance(f.filter_id, f.options) for f in data.filters] if data.filters is not None else None
        template = store.update_template(
            template_id=template_id,
            name=data.name,
            filters=filters,
            description=data.description,
        )
        return _pp_template_to_response(template)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to update postprocessing template: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.delete("/api/v1/postprocessing-templates/{template_id}", status_code=204, tags=["Postprocessing Templates"])
async def delete_pp_template(
    template_id: str,
    _auth: AuthRequired,
    store: PostprocessingTemplateStore = Depends(get_pp_template_store),
    stream_store: PictureSourceStore = Depends(get_picture_source_store),
):
    """Delete a postprocessing template."""
    try:
        # Check if any picture source references this template
        if store.is_referenced_by(template_id, stream_store):
            raise HTTPException(
                status_code=409,
                detail="Cannot delete postprocessing template: it is referenced by one or more picture sources. "
                       "Please reassign those streams before deleting.",
            )
        store.delete_template(template_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 postprocessing template: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/api/v1/postprocessing-templates/{template_id}/test", response_model=TemplateTestResponse, tags=["Postprocessing Templates"])
async def test_pp_template(
    template_id: str,
    test_request: PPTemplateTestRequest,
    _auth: AuthRequired,
    pp_store: PostprocessingTemplateStore = Depends(get_pp_template_store),
    stream_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),
):
    """Test a postprocessing template by capturing from a source stream and applying filters."""
    stream = None
    try:
        # Get the PP template
        try:
            pp_template = pp_store.get_template(template_id)
        except ValueError as e:
            raise HTTPException(status_code=404, detail=str(e))

        # Resolve source stream chain to get the raw stream
        try:
            chain = stream_store.resolve_stream_chain(test_request.source_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: load directly
            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()

            logger.info(f"Starting {test_request.capture_duration}s PP template test for {template_id} using stream {test_request.source_stream_id}")

            frame_count = 0
            total_capture_time = 0.0
            last_frame = None

            start_time = time.perf_counter()
            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 pp_template.filters:
            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)

        # 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 Exception as e:
        logger.error(f"Postprocessing template test failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if stream:
            try:
                stream.cleanup()
            except Exception:
                pass