core/manager.py

"""High-level inference execution and observability."""

from __future__ import annotations

from pathlib import Path
from typing import Any, Callable, Literal, Protocol, TypeVar

from config import Settings
from core.errors import GatewayError
from core.executor import InferenceCommand, execute_inference
from core.loader import ModelLoader
from core.queue import InferenceQueue
from core.runtime import gpu_available, memory_stats, zerogpu_enabled
from core.schemas import (
    EpisodeAssets,
    ImageAsset,
    ImageRequest,
    MusicAsset,
    MusicRequest,
    SFXAsset,
    SFXRequest,
    SpeechAsset,
    TTSRequest,
    TranscriptionAsset,
    VideoAsset,
    VideoRequest,
    WorkflowRequest,
)
from utils.files import OutputManager
from utils.validation import validate_image_dimensions


T = TypeVar("T")


class WorkflowRunner(Protocol):
    """Structural contract that avoids coupling the manager to workflow internals."""

    async def run(
        self, payload: WorkflowRequest, request_id: str
    ) -> dict[str, Path]:
        ...


class TaskManager:
    """Submit inference to one queue and manage model lifecycle consistently."""

    def __init__(self, settings: Settings, loader: ModelLoader) -> None:
        self.settings = settings
        self.loader = loader
        self.queue = InferenceQueue(settings.max_queue, settings.job_timeout_seconds)

    async def start(self) -> None:
        await self.queue.start()

    async def stop(self) -> None:
        await self.queue.stop()

    async def run(
        self,
        command: InferenceCommand,
    ) -> Any:
        """Queue one model action."""
        return await self.queue.submit(
            command.request_id,
            command.model_name,
            lambda: self.invoke_direct(command),
        )

    async def run_exclusive(
        self, request_id: str, label: str, action: Callable[[], T]
    ) -> T:
        """Queue a compound operation as one non-interleavable job."""
        return await self.queue.submit(request_id, label, action)

    @staticmethod
    def invoke_direct(command: InferenceCommand) -> Any:
        """Invoke a serializable command; workflows reuse this method."""
        return execute_inference(command)


class AIService:
    """Shared AI capability layer used by every external interface."""

    model_names = ("flux", "wan22", "kokoro", "musicgen", "whisper", "sfx")
    rest_endpoints = (
        "/api/image",
        "/api/video",
        "/tts",
        "/music",
        "/sfx",
        "/transcribe",
        "/workflow",
        "/health",
    )

    def __init__(
        self,
        settings: Settings,
        tasks: TaskManager,
        loader: ModelLoader,
        outputs: OutputManager,
        workflows: WorkflowRunner,
    ) -> None:
        self.settings = settings
        self.tasks = tasks
        self.loader = loader
        self.outputs = outputs
        self.workflows = workflows

    async def generate_image(self, payload: ImageRequest, request_id: str) -> ImageAsset:
        """Validate, queue, and persist one FLUX image."""
        width = payload.width if payload.width is not None else self.settings.image_width
        height = payload.height if payload.height is not None else self.settings.image_height
        steps = payload.steps if payload.steps is not None else self.settings.image_steps
        guidance_scale = (
            payload.guidance_scale
            if payload.guidance_scale is not None
            else self.settings.flux_guidance_scale
        )
        validate_image_dimensions(width, height)
        target = self.outputs.allocate("images")
        try:
            await self.tasks.run(
                InferenceCommand(
                    model_name="flux",
                    method_name="generate",
                    arguments={
                        "prompt": payload.prompt,
                        "output_path": target,
                        "width": width,
                        "height": height,
                        "steps": steps,
                        "seed": payload.seed,
                        "guidance_scale": guidance_scale,
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_flux_duration,
                ),
            )
        except Exception:
            self.outputs.remove(target)
            raise
        return ImageAsset(image=self.outputs.public_path(target))

    async def generate_video(self, payload: VideoRequest, request_id: str) -> VideoAsset:
        """Validate, queue, and persist a WAN image-to-video result."""
        frames = payload.frames if payload.frames is not None else self.settings.video_frames
        fps = payload.fps if payload.fps is not None else self.settings.video_fps
        steps = payload.steps if payload.steps is not None else self.settings.wan_steps
        guidance_scale = (
            payload.guidance_scale
            if payload.guidance_scale is not None
            else self.settings.wan_guidance_scale
        )
        if (frames - 1) % 4:
            raise GatewayError("frames must equal 4k + 1", code="invalid_frame_count")
        source = self.outputs.resolve_public_input(payload.image, categories={"images"})
        target = self.outputs.allocate("videos")
        try:
            await self.tasks.run(
                InferenceCommand(
                    model_name="wan",
                    method_name="generate",
                    arguments={
                        "image_path": source,
                        "prompt": payload.prompt,
                        "negative_prompt": payload.negative_prompt,
                        "output_path": target,
                        "steps": steps,
                        "frames": frames,
                        "fps": fps,
                        "seed": payload.seed,
                        "guidance_scale": guidance_scale,
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_wan_duration,
                    gpu_size="xlarge",
                ),
            )
        except Exception:
            self.outputs.remove(target)
            raise
        return VideoAsset(video=self.outputs.public_path(target))

    async def generate_speech(self, payload: TTSRequest, request_id: str) -> SpeechAsset:
        """Validate, queue, and persist Kokoro speech."""
        voice = payload.voice or self.settings.kokoro_default_voice
        target = self.outputs.allocate("audio")
        try:
            await self.tasks.run(
                InferenceCommand(
                    model_name="kokoro",
                    method_name="synthesize",
                    arguments={
                        "text": payload.text,
                        "voice": voice,
                        "speed": payload.speed,
                        "output_path": target,
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_kokoro_duration,
                ),
            )
        except Exception:
            self.outputs.remove(target)
            raise
        return SpeechAsset(voice=self.outputs.public_path(target))

    async def generate_music(self, payload: MusicRequest, request_id: str) -> MusicAsset:
        """Validate, queue, and persist MusicGen output."""
        duration = (
            payload.duration if payload.duration is not None else self.settings.music_duration
        )
        guidance_scale = (
            payload.guidance_scale
            if payload.guidance_scale is not None
            else self.settings.music_guidance_scale
        )
        target = self.outputs.allocate("music")
        try:
            await self.tasks.run(
                InferenceCommand(
                    model_name="musicgen",
                    method_name="generate",
                    arguments={
                        "prompt": payload.prompt,
                        "duration": duration,
                        "guidance_scale": guidance_scale,
                        "seed": payload.seed,
                        "output_path": target,
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_musicgen_duration,
                ),
            )
        except Exception:
            self.outputs.remove(target)
            raise
        return MusicAsset(music=self.outputs.public_path(target))

    async def generate_sfx(self, payload: SFXRequest, request_id: str) -> SFXAsset:
        """Validate, queue, and persist an AudioLDM2 sound effect."""
        duration = (
            payload.duration if payload.duration is not None else self.settings.sfx_duration
        )
        steps = payload.steps if payload.steps is not None else self.settings.sfx_steps
        target = self.outputs.allocate("audio")
        try:
            await self.tasks.run(
                InferenceCommand(
                    model_name="sfx",
                    method_name="generate",
                    arguments={
                        "prompt": payload.prompt,
                        "duration": duration,
                        "steps": steps,
                        "seed": payload.seed,
                        "output_path": target,
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_sfx_duration,
                ),
            )
        except Exception:
            self.outputs.remove(target)
            raise
        return SFXAsset(sfx=self.outputs.public_path(target))

    async def transcribe_source(
        self,
        source: Path,
        request_id: str,
        *,
        language: str | None = None,
        task: Literal["transcribe", "translate"] = "transcribe",
    ) -> TranscriptionAsset:
        """Queue faster-whisper for a trusted local source file."""
        subtitle = self.outputs.allocate("subtitles")
        try:
            result = await self.tasks.run(
                InferenceCommand(
                    model_name="whisper",
                    method_name="transcribe",
                    arguments={
                        "source": source,
                        "subtitle_path": subtitle,
                        "language": language,
                        "task": task,
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_whisper_duration,
                ),
            )
        except Exception:
            self.outputs.remove(subtitle)
            raise
        return TranscriptionAsset(
            text=result["text"],
            subtitle=self.outputs.public_path(subtitle),
            language=result["language"],
        )

    async def transcribe_output(
        self,
        audio: str,
        request_id: str,
        *,
        language: str | None = None,
        task: Literal["transcribe", "translate"] = "transcribe",
    ) -> TranscriptionAsset:
        """Resolve and transcribe a gateway-generated audio or video asset."""
        source = self.outputs.resolve_public_input(
            audio, categories={"audio", "music", "videos"}
        )
        return await self.transcribe_source(
            source, request_id, language=language, task=task
        )

    async def create_episode(
        self, payload: WorkflowRequest, request_id: str
    ) -> EpisodeAssets:
        """Run the full workflow as one exclusive queue job."""
        width = payload.width if payload.width is not None else self.settings.image_width
        height = payload.height if payload.height is not None else self.settings.image_height
        validate_image_dimensions(width, height)
        assets = await self.workflows.run(payload, request_id)
        return EpisodeAssets(
            image=self.outputs.public_path(assets["image"]),
            video=self.outputs.public_path(assets["video"]),
            voice=self.outputs.public_path(assets["voice"]),
            music=self.outputs.public_path(assets["music"]),
            subtitle=self.outputs.public_path(assets["subtitle"]),
        )

    def health(self) -> dict[str, object]:
        """Return a no-load health snapshot for REST and MCP."""
        return {
            "status": "healthy",
            "gpu": "available" if gpu_available() else "unavailable",
            "queue": self.tasks.queue.depth,
            "models_loaded": self._loaded_models(),
        }

    def list_models(self) -> dict[str, list[str]]:
        """Return every discoverable gateway model capability."""
        return {"models": list(self.model_names)}

    def server_info(self) -> dict[str, object]:
        """Return operational metadata without loading a model."""
        return {
            "version": self.settings.app_version,
            "device": "zerogpu" if zerogpu_enabled() else self.loader.device,
            "gpu": "available" if gpu_available() else "unavailable",
            "memory": memory_stats(),
            "available_endpoints": [*self.rest_endpoints, "/mcp"],
            "loaded_models": self._loaded_models(),
            "queue": {
                "depth": self.tasks.queue.depth,
                "capacity": self.tasks.queue.capacity,
                "active": self.tasks.queue.active,
            },
        }

    def _loaded_models(self) -> list[str]:
        """Expose stable public model names rather than adapter implementation names."""
        return [
            "wan22" if model_name == "wan" else model_name
            for model_name in self.loader.loaded_models
        ]