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

"""Audio capture template and engine routes."""

import asyncio
import json
import secrets

from fastapi import APIRouter, HTTPException, Depends, Query
from starlette.websockets import WebSocket, WebSocketDisconnect

from wled_controller.api.auth import AuthRequired
from wled_controller.api.dependencies import get_audio_template_store, get_audio_source_store, get_processor_manager
from wled_controller.api.schemas.audio_templates import (
    AudioEngineInfo,
    AudioEngineListResponse,
    AudioTemplateCreate,
    AudioTemplateListResponse,
    AudioTemplateResponse,
    AudioTemplateUpdate,
)
from wled_controller.config import get_config
from wled_controller.core.audio.factory import AudioEngineRegistry
from wled_controller.storage.audio_template_store import AudioTemplateStore
from wled_controller.storage.audio_source_store import AudioSourceStore
from wled_controller.utils import get_logger

logger = get_logger(__name__)

router = APIRouter()


# ===== AUDIO TEMPLATE ENDPOINTS =====

@router.get("/api/v1/audio-templates", response_model=AudioTemplateListResponse, tags=["Audio Templates"])
async def list_audio_templates(
    _auth: AuthRequired,
    store: AudioTemplateStore = Depends(get_audio_template_store),
):
    """List all audio capture templates."""
    try:
        templates = store.get_all_templates()
        responses = [
            AudioTemplateResponse(
                id=t.id, name=t.name, engine_type=t.engine_type,
                engine_config=t.engine_config, created_at=t.created_at,
                updated_at=t.updated_at, description=t.description,
            )
            for t in templates
        ]
        return AudioTemplateListResponse(templates=responses, count=len(responses))
    except Exception as e:
        logger.error(f"Failed to list audio templates: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/api/v1/audio-templates", response_model=AudioTemplateResponse, tags=["Audio Templates"], status_code=201)
async def create_audio_template(
    data: AudioTemplateCreate,
    _auth: AuthRequired,
    store: AudioTemplateStore = Depends(get_audio_template_store),
):
    """Create a new audio capture template."""
    try:
        template = store.create_template(
            name=data.name, engine_type=data.engine_type,
            engine_config=data.engine_config, description=data.description,
        )
        return AudioTemplateResponse(
            id=template.id, name=template.name, engine_type=template.engine_type,
            engine_config=template.engine_config, created_at=template.created_at,
            updated_at=template.updated_at, description=template.description,
        )
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to create audio template: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/api/v1/audio-templates/{template_id}", response_model=AudioTemplateResponse, tags=["Audio Templates"])
async def get_audio_template(
    template_id: str,
    _auth: AuthRequired,
    store: AudioTemplateStore = Depends(get_audio_template_store),
):
    """Get audio template by ID."""
    try:
        t = store.get_template(template_id)
    except ValueError:
        raise HTTPException(status_code=404, detail=f"Audio template {template_id} not found")
    return AudioTemplateResponse(
        id=t.id, name=t.name, engine_type=t.engine_type,
        engine_config=t.engine_config, created_at=t.created_at,
        updated_at=t.updated_at, description=t.description,
    )


@router.put("/api/v1/audio-templates/{template_id}", response_model=AudioTemplateResponse, tags=["Audio Templates"])
async def update_audio_template(
    template_id: str,
    data: AudioTemplateUpdate,
    _auth: AuthRequired,
    store: AudioTemplateStore = Depends(get_audio_template_store),
):
    """Update an audio template."""
    try:
        t = store.update_template(
            template_id=template_id, name=data.name,
            engine_type=data.engine_type, engine_config=data.engine_config,
            description=data.description,
        )
        return AudioTemplateResponse(
            id=t.id, name=t.name, engine_type=t.engine_type,
            engine_config=t.engine_config, created_at=t.created_at,
            updated_at=t.updated_at, description=t.description,
        )
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Failed to update audio template: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.delete("/api/v1/audio-templates/{template_id}", status_code=204, tags=["Audio Templates"])
async def delete_audio_template(
    template_id: str,
    _auth: AuthRequired,
    store: AudioTemplateStore = Depends(get_audio_template_store),
    audio_source_store: AudioSourceStore = Depends(get_audio_source_store),
):
    """Delete an audio template."""
    try:
        store.delete_template(template_id, audio_source_store=audio_source_store)
    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 audio template: {e}")
        raise HTTPException(status_code=500, detail=str(e))


# ===== AUDIO ENGINE ENDPOINTS =====

@router.get("/api/v1/audio-engines", response_model=AudioEngineListResponse, tags=["Audio Templates"])
async def list_audio_engines(_auth: AuthRequired):
    """List all registered audio capture engines."""
    try:
        available_set = set(AudioEngineRegistry.get_available_engines())
        all_engines = AudioEngineRegistry.get_all_engines()

        engines = []
        for engine_type, engine_class in all_engines.items():
            engines.append(
                AudioEngineInfo(
                    type=engine_type,
                    name=engine_type.upper(),
                    default_config=engine_class.get_default_config(),
                    available=(engine_type in available_set),
                )
            )

        return AudioEngineListResponse(engines=engines, count=len(engines))
    except Exception as e:
        logger.error(f"Failed to list audio engines: {e}")
        raise HTTPException(status_code=500, detail=str(e))


# ===== REAL-TIME AUDIO TEMPLATE TEST WEBSOCKET =====


@router.websocket("/api/v1/audio-templates/{template_id}/test/ws")
async def test_audio_template_ws(
    websocket: WebSocket,
    template_id: str,
    token: str = Query(""),
    device_index: int = Query(-1),
    is_loopback: int = Query(1),
):
    """WebSocket for real-time audio spectrum test of a template with a chosen device.

    Auth via ?token=<api_key>. Device specified via ?device_index=N&is_loopback=0|1.
    Streams AudioAnalysis snapshots as JSON at ~20 Hz.
    """
    # Authenticate
    authenticated = False
    cfg = get_config()
    if token and cfg.auth.api_keys:
        for _label, api_key in cfg.auth.api_keys.items():
            if secrets.compare_digest(token, api_key):
                authenticated = True
                break

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

    # Resolve template
    store = get_audio_template_store()
    try:
        template = store.get_template(template_id)
    except ValueError:
        await websocket.close(code=4004, reason="Template not found")
        return

    # Acquire shared audio stream
    manager = get_processor_manager()
    audio_mgr = manager.audio_capture_manager
    loopback = is_loopback != 0

    try:
        stream = audio_mgr.acquire(device_index, loopback, template.engine_type, template.engine_config)
    except RuntimeError as e:
        await websocket.close(code=4003, reason=str(e))
        return

    await websocket.accept()
    logger.info(f"Audio template test WS connected: template={template_id} device={device_index} loopback={loopback}")

    last_ts = 0.0
    try:
        while True:
            analysis = stream.get_latest_analysis()
            if analysis is not None and analysis.timestamp != last_ts:
                last_ts = analysis.timestamp
                await websocket.send_json({
                    "spectrum": analysis.spectrum.tolist(),
                    "rms": round(analysis.rms, 4),
                    "peak": round(analysis.peak, 4),
                    "beat": analysis.beat,
                    "beat_intensity": round(analysis.beat_intensity, 4),
                })
            await asyncio.sleep(0.05)
    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error(f"Audio template test WS error: {e}")
    finally:
        audio_mgr.release(device_index, loopback, template.engine_type)
        logger.info(f"Audio template test WS disconnected: template={template_id}")