app/templates/executor.py

from __future__ import annotations

import os
import time
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Literal

from app.core.exceptions import (
    MediaAPIError,
    TemplateExecutionError,
    TemplateValidationError,
)
from app.core.logger import get_logger
from app.core.response import SuccessResponse
from app.models.media import InputMedia, MediaSource, OperationResult, ResolvedRequest
from app.operations.compress import compress_video, normalize_audio, normalize_video
from app.operations.concat import (
    concat_audio,
    concat_video,
    image_sequence,
    image_slideshow,
    image_to_video,
)
from app.operations.convert import convert_audio, convert_image, convert_video
from app.operations.crop import crop_image, crop_video
from app.operations.extract_audio import (
    extract_audio,
    fade_audio,
    mute_video,
    noise_reduction,
    remove_audio,
    remove_silence,
    replace_audio,
    set_volume,
)
from app.operations.merge import merge_audio, merge_video
from app.operations.resize import pad_video, resize_image, resize_video, scale_video
from app.operations.rotate import (
    change_bitrate,
    change_fps,
    change_speed,
    reverse_video,
    rotate_video,
    slow_motion,
)
from app.operations.subtitles import burn_subtitles, soft_subtitles
from app.operations.thumbnails import (
    blur_video,
    denoise_video,
    extract_frames,
    generate_gif,
    sharpen_video,
    thumbnail,
)
from app.operations.trim import trim_audio, trim_video
from app.operations.watermark import (
    overlay_image,
    overlay_video,
    watermark_image,
    watermark_video,
)
from app.services.media_service import MediaProcessor, Operation
from app.templates.models import PreparedTemplate, ResolvedPipelineStep
from app.templates.registry import TemplateRegistry

logger = get_logger(__name__)

OperationKind = Literal["ffmpeg", "whisper", "passthrough"]
InputMode = Literal["current", "all"]


@dataclass(frozen=True, slots=True)
class OperationBinding:
    """Binding from a safe YAML operation name to shared application behavior."""

    kind: OperationKind
    handler: Operation | None = None
    input_mode: InputMode = "current"
    whisper_task: Literal["transcribe", "translate"] = "transcribe"
    aliases: tuple[tuple[str, str], ...] = ()


def _ffmpeg(
    handler: Operation,
    *,
    input_mode: InputMode = "current",
    aliases: tuple[tuple[str, str], ...] = (),
) -> OperationBinding:
    return OperationBinding(kind="ffmpeg", handler=handler, input_mode=input_mode, aliases=aliases)


OPERATION_BINDINGS: Mapping[str, OperationBinding] = MappingProxyType(
    {
        "compress": _ffmpeg(compress_video),
        "compress_video": _ffmpeg(compress_video),
        "resize": _ffmpeg(resize_video),
        "resize_video": _ffmpeg(resize_video),
        "crop": _ffmpeg(crop_video),
        "crop_video": _ffmpeg(crop_video),
        "trim": _ffmpeg(trim_video),
        "trim_video": _ffmpeg(trim_video),
        "rotate": _ffmpeg(rotate_video),
        "reverse": _ffmpeg(reverse_video),
        "merge": _ffmpeg(merge_video, input_mode="all"),
        "merge_videos": _ffmpeg(merge_video, input_mode="all"),
        "concat": _ffmpeg(concat_video, input_mode="all"),
        "concat_videos": _ffmpeg(concat_video, input_mode="all"),
        "convert": _ffmpeg(convert_video),
        "convert_video": _ffmpeg(convert_video),
        "overlay": _ffmpeg(overlay_video, input_mode="all"),
        "overlay_video": _ffmpeg(overlay_video, input_mode="all"),
        "watermark": _ffmpeg(watermark_video, input_mode="all"),
        "watermark_video": _ffmpeg(watermark_video, input_mode="all"),
        "extract_frames": _ffmpeg(extract_frames),
        "generate_gif": _ffmpeg(generate_gif),
        "thumbnail": _ffmpeg(thumbnail),
        "replace_audio": _ffmpeg(replace_audio, input_mode="all"),
        "remove_audio": _ffmpeg(remove_audio),
        "mute": _ffmpeg(mute_video),
        "speed": _ffmpeg(change_speed),
        "slow_motion": _ffmpeg(slow_motion),
        "fps": _ffmpeg(change_fps, aliases=(("value", "fps"),)),
        "bitrate": _ffmpeg(change_bitrate, aliases=(("value", "bitrate"),)),
        "burn_subtitles": _ffmpeg(burn_subtitles, input_mode="all"),
        "soft_subtitles": _ffmpeg(soft_subtitles, input_mode="all"),
        "scale": _ffmpeg(scale_video),
        "pad": _ffmpeg(pad_video),
        "blur": _ffmpeg(blur_video),
        "sharpen": _ffmpeg(sharpen_video),
        "denoise": _ffmpeg(denoise_video),
        "normalize_video": _ffmpeg(normalize_video),
        "extract_audio": _ffmpeg(extract_audio),
        "convert_audio": _ffmpeg(convert_audio),
        "normalize_audio": _ffmpeg(normalize_audio),
        "trim_audio": _ffmpeg(trim_audio),
        "merge_audio": _ffmpeg(merge_audio, input_mode="all"),
        "concat_audio": _ffmpeg(concat_audio, input_mode="all"),
        "fade_audio": _ffmpeg(fade_audio),
        "volume": _ffmpeg(set_volume),
        "remove_silence": _ffmpeg(remove_silence),
        "noise_reduction": _ffmpeg(noise_reduction),
        "resize_image": _ffmpeg(resize_image),
        "crop_image": _ffmpeg(crop_image),
        "convert_image": _ffmpeg(convert_image),
        "slideshow": _ffmpeg(image_slideshow, input_mode="all"),
        "image_sequence": _ffmpeg(image_sequence, input_mode="all"),
        "image_to_video": _ffmpeg(image_to_video),
        "watermark_image": _ffmpeg(watermark_image, input_mode="all"),
        "overlay_image": _ffmpeg(overlay_image, input_mode="all"),
        "transcribe": OperationBinding(kind="whisper", whisper_task="transcribe"),
        "translate": OperationBinding(kind="whisper", whisper_task="translate"),
        "download": OperationBinding(kind="passthrough"),
    }
)


class OperationExecutor:
    """Executes allow-listed YAML operations through existing implementations."""

    def __init__(self, processor: MediaProcessor) -> None:
        self.processor = processor

    @property
    def supported_operations(self) -> set[str]:
        """Return operation names that are safe for template YAML."""
        return set(OPERATION_BINDINGS)

    def input_mode(self, operation: str) -> InputMode:
        """Return the default input selection mode for an operation."""
        return self._binding(operation).input_mode

    async def execute(
        self,
        operation: str,
        inputs: list[InputMedia],
        parameters: dict[str, object],
        output_dir: Path,
    ) -> OperationResult:
        """Execute one pipeline step without publishing its intermediate output."""
        binding = self._binding(operation)
        normalized_parameters = dict(parameters)
        for source, destination in binding.aliases:
            if source in normalized_parameters and destination not in normalized_parameters:
                normalized_parameters[destination] = normalized_parameters[source]
        if binding.kind == "passthrough":
            if not inputs:
                raise TemplateExecutionError("Download step requires an input")
            media = inputs[0]
            return OperationResult(
                path=media.temp_path,
                filename=media.filename,
                mime_type=media.mime_type,
                metadata={"operation": "download", **media.metadata},
            )
        if binding.kind == "whisper":
            normalized_parameters["task"] = binding.whisper_task
            return await self.processor.transcribe_result(inputs, normalized_parameters, output_dir)
        if binding.handler is None:  # pragma: no cover - guarded by static bindings
            raise TemplateExecutionError("Template operation has no implementation")
        return await binding.handler(
            self.processor.ffmpeg, inputs, normalized_parameters, output_dir
        )

    @staticmethod
    def _binding(operation: str) -> OperationBinding:
        binding = OPERATION_BINDINGS.get(operation)
        if binding is None:
            raise TemplateValidationError(f"Unsupported template operation '{operation}'")
        return binding


class TemplateExecutor:
    """Runs validated template pipelines over normalized InputMedia instances."""

    def __init__(
        self,
        registry: TemplateRegistry,
        operation_executor: OperationExecutor,
        processor: MediaProcessor,
    ) -> None:
        self.registry = registry
        self.operation_executor = operation_executor
        self.processor = processor

    async def execute_request(self, resolved: ResolvedRequest) -> SuccessResponse:
        """Execute template controls parsed by the shared InputResolver."""
        reference = resolved.params.get("template")
        if not isinstance(reference, str) or not reference.strip():
            raise TemplateValidationError("A non-empty 'template' reference is required")
        parameters = resolved.params.get("parameters", {})
        if not isinstance(parameters, dict):
            raise TemplateValidationError("Template 'parameters' must be an object")
        return await self.execute(resolved, reference, parameters)

    async def execute(
        self,
        resolved: ResolvedRequest,
        template_reference: str,
        parameters: dict[str, object] | None = None,
    ) -> SuccessResponse:
        """Execute a versioned template and publish only its final artifact."""
        started = time.monotonic()
        cpu_started = time.process_time()
        prepared: PreparedTemplate | None = None
        operations: list[str] = []
        try:
            prepared = self.registry.prepare(template_reference, parameters)
            workspace = await self.processor.cleanup.create_workspace(resolved.request_id)
            initial_metadata = await self.processor.probe_inputs(resolved.inputs)
            originals = list(resolved.inputs)
            current = originals[0]
            artifacts: dict[str, InputMedia] = {}
            for index, step in enumerate(prepared.pipeline):
                if not step.enabled:
                    continue
                selected = self._select_inputs(step, current, originals, artifacts)
                step_dir = workspace.outputs / f"{index + 1:03d}_{step.operation}"
                result = await self.operation_executor.execute(
                    step.operation, selected, step.parameters, step_dir
                )
                current = self._result_media(result)
                if self._probeable(current):
                    await self.processor.probe_inputs([current])
                if step.save_as:
                    artifacts[step.save_as] = current
                operations.append(step.operation)
            if not operations:
                raise TemplateExecutionError("All template pipeline operations were disabled")
            self._validate_output(prepared, current)
            result = OperationResult(
                path=current.temp_path,
                filename=prepared.output.filename or current.filename,
                mime_type=current.mime_type,
                metadata=current.metadata,
            )
            response = await self.processor.finish_result(
                resolved,
                f"template.{prepared.definition.id}",
                result,
                started,
                initial_metadata,
                extra_metadata={
                    "template": {
                        "id": prepared.definition.id,
                        "version": prepared.definition.version,
                        "reference": (f"{prepared.definition.id}@{prepared.definition.version}"),
                        "category": prepared.definition.category,
                    },
                    "parameters": prepared.parameters,
                    "operations": operations,
                },
            )
            logger.info(
                "template execution completed",
                extra=self._log_data(
                    prepared,
                    resolved,
                    operations,
                    started,
                    cpu_started,
                    response.metadata.get("output_size", 0),
                    template_reference,
                    parameters,
                ),
            )
            return response
        except MediaAPIError as exc:
            logger.warning(
                "template execution failed",
                extra={
                    **self._log_data(
                        prepared,
                        resolved,
                        operations,
                        started,
                        cpu_started,
                        0,
                        template_reference,
                        parameters,
                    ),
                    "error_code": exc.code,
                    "error": exc.message,
                },
            )
            raise
        except Exception:
            logger.exception(
                "unexpected template execution error",
                extra=self._log_data(
                    prepared,
                    resolved,
                    operations,
                    started,
                    cpu_started,
                    0,
                    template_reference,
                    parameters,
                ),
            )
            raise
        finally:
            try:
                await self.processor.cleanup.complete(resolved.request_id)
            except Exception:
                logger.exception(
                    "template workspace completion failed",
                    extra={"request_id": resolved.request_id},
                )

    def _select_inputs(
        self,
        step: ResolvedPipelineStep,
        current: InputMedia,
        originals: list[InputMedia],
        artifacts: dict[str, InputMedia],
    ) -> list[InputMedia]:
        if step.inputs is None:
            if self.operation_executor.input_mode(step.operation) == "all":
                return [current, *originals[1:]]
            return [current]
        selected: list[InputMedia] = []
        for selector in step.inputs:
            if selector == "current":
                selected.append(current)
            elif selector == "original":
                selected.append(originals[0])
            elif selector == "originals":
                selected.extend(originals)
            elif selector.startswith("original:"):
                index = int(selector.split(":", 1)[1])
                try:
                    selected.append(originals[index])
                except IndexError as exc:
                    raise TemplateExecutionError(
                        f"Template requires original input index {index}"
                    ) from exc
            else:
                name = selector.split(":", 1)[1]
                try:
                    selected.append(artifacts[name])
                except KeyError as exc:  # pragma: no cover - definition validation guards order
                    raise TemplateExecutionError(
                        f"Template artifact '{name}' is unavailable"
                    ) from exc
        if not selected:
            raise TemplateExecutionError("Template operation selected no inputs")
        return selected

    def _result_media(self, result: OperationResult) -> InputMedia:
        if result.path is None or not result.path.is_file():
            raise TemplateExecutionError("A template operation did not produce a file artifact")
        filename = result.filename or result.path.name
        mime_type = result.mime_type or self.processor.validator.infer_mime(result.path)
        return InputMedia(
            source=MediaSource.LOCAL_PATH,
            filename=filename,
            mime_type=mime_type,
            temp_path=result.path,
            size=result.path.stat().st_size,
            metadata=result.metadata,
        )

    @staticmethod
    def _probeable(media: InputMedia) -> bool:
        return media.mime_type.startswith(("video/", "audio/", "image/"))

    @staticmethod
    def _validate_output(prepared: PreparedTemplate, media: InputMedia) -> None:
        if prepared.output.filename:
            filename = prepared.output.filename
            if Path(filename).name != filename or len(filename) > 255:
                raise TemplateValidationError("Template output filename must be a safe basename")
        expected = prepared.output.format.lower().lstrip(".")
        if expected == "source":
            return
        actual = media.temp_path.suffix.lower().lstrip(".")
        aliases = {"jpeg": "jpg", "m4a": "m4a"}
        if aliases.get(actual, actual) != aliases.get(expected, expected):
            raise TemplateExecutionError(
                "Template output did not match its declared format",
                details={"expected": expected, "actual": actual},
            )

    @staticmethod
    def _log_data(
        prepared: PreparedTemplate | None,
        resolved: ResolvedRequest,
        operations: list[str],
        started: float,
        cpu_started: float,
        output_size: object,
        template_reference: str,
        supplied_parameters: dict[str, object] | None,
    ) -> dict[str, object]:
        memory: int | None = None
        cpu_percent: float | None = None
        try:
            import psutil

            memory = psutil.Process(os.getpid()).memory_info().rss
            cpu_percent = psutil.cpu_percent(interval=None)
        except ImportError:
            pass
        return {
            "template_id": prepared.definition.id if prepared else template_reference,
            "template_version": prepared.definition.version if prepared else None,
            "template_reference": template_reference,
            "parameters": prepared.parameters if prepared else supplied_parameters or {},
            "request_id": resolved.request_id,
            "operations": operations,
            "duration": round(time.monotonic() - started, 4),
            "cpu_time": round(time.process_time() - cpu_started, 6),
            "cpu_percent": cpu_percent,
            "memory_bytes": memory,
            "output_size": output_size,
        }