renderer/core/render_engine.py

from __future__ import annotations

import mimetypes
import shutil
import time
from pathlib import Path

from renderer.audio import AudioMixer
from renderer.core.config import Settings
from renderer.core.ingest import AssetIngestor
from renderer.core.models import AIReelsRequest, RenderRequest, RenderResult
from renderer.core.utils import cleanup_directory, temp_workdir
from renderer.exports import ExportManager
from renderer.ffmpeg.assets import AssetProbe
from renderer.ffmpeg.command import FFmpegCommand
from renderer.ffmpeg.normalize import Normalizer
from renderer.ffmpeg.runner import FFmpegRunner
from renderer.scenes import Timeline
from renderer.subtitles import SubtitleGenerator
from renderer.templates import PlatformProfile, get_creative_style, get_platform_profile, scene_effect_filter
from renderer.transcription import WhisperTranscriber
from renderer.transitions import TransitionBuilder


class RenderEngine:
    def __init__(self, settings: Settings | None = None, log=None) -> None:
        self.settings = settings or Settings()
        self.settings.ensure_dirs()
        self._commands: list[list[str]] = []
        self._logs: list[str] = []
        self.runner = FFmpegRunner(self.settings.ffmpeg_timeout_seconds, log=log, on_command=self._record_command)
        self.ingest = AssetIngestor(self.settings)
        self.assets = AssetProbe(self.settings.metadata_cache)
        self.normalizer = Normalizer(self.settings, self.runner)
        self.subtitles = SubtitleGenerator()
        self.transcriber = WhisperTranscriber(self.settings)
        self.transitions = TransitionBuilder()
        self.audio = AudioMixer(self.runner)
        self.exports = ExportManager(self.settings.exports_dir)

    def render(self, request: RenderRequest, job_id: str) -> RenderResult:
        self._commands = []
        self._logs = []
        started = time.time()
        timeline = Timeline(request.scenes)
        with temp_workdir(self.settings.temp_dir, job_id) as work:
            workdir = Path(work)
            resolved = self.ingest.resolve_render_request(request, workdir)
            timeline = Timeline(resolved.scenes)
            prepared = self._prepare_scene_media(resolved, workdir)
            subtitles = self._write_subtitles(resolved, timeline, workdir)
            video = self._compose_video(prepared, resolved, subtitles, workdir)
            mixed = self.audio.mix(
                video,
                resolved.background_music,
                resolved.voiceover,
                workdir / "mixed.mp4",
                normalize=resolved.audio_normalize,
                duration=timeline.total_duration,
                music_volume=resolved.music_volume,
                music_fade_in=resolved.music_fade_in,
                music_fade_out=resolved.music_fade_out,
                music_loop=resolved.music_loop,
                music_start=resolved.music_start,
                music_ducking=resolved.music_ducking,
                voice_volume=resolved.voice_volume,
            )
            mixed = self._apply_watermark(mixed, resolved, workdir)
            optimized = self._optimize_for_platform(mixed, resolved, workdir)
            output = self.exports.save(optimized, job_id, request.output_name)
            cleanup_directory(workdir, keep={mixed})
        profile = self._profile(request)
        metrics = {
            "render_time_seconds": round(time.time() - started, 3),
            "output_size_bytes": output.stat().st_size,
            "scene_count": len(request.scenes),
            "creative_style": request.creative_style or request.metadata.get("creative_style"),
            "scene_effects": [scene.effect for scene in request.scenes if scene.effect],
            "platform": profile.metadata(),
        }
        return RenderResult(output_path=output, commands=list(self._commands), metrics=metrics, logs=list(self._logs))

    def ai_reels(self, request: AIReelsRequest, job_id: str) -> RenderResult:
        if not request.voiceover:
            raise ValueError("AI reels v1 requires a provided voiceover path; TTS is pluggable but not bundled.")
        if not request.assets:
            raise ValueError("AI reels rendering requires at least one visual asset.")
        with temp_workdir(self.settings.temp_dir, job_id) as work:
            workdir = Path(work)
            resolved = self.ingest.resolve_ai_reels_request(request, workdir)
            style = get_creative_style(resolved.creative_style)
            voice_meta = self.assets.probe(resolved.voiceover)
            duration = max(voice_meta.duration, len(resolved.script.split()) * 0.35, style.scene_duration * len(resolved.assets), 3.0)
            per_scene = duration / max(1, len(resolved.assets))
            captions = _split_script(resolved.script, len(resolved.assets))
            scenes = [
                {
                    "start": round(idx * per_scene, 3),
                    "duration": round(per_scene, 3),
                    "media": asset,
                    "caption": captions[idx] if idx < len(captions) else "",
                    "transition": style.transition_sequence[idx % len(style.transition_sequence)],
                    "effect": style.scene_effect_sequence[idx % len(style.scene_effect_sequence)],
                }
                for idx, asset in enumerate(resolved.assets)
            ]
            render_request = RenderRequest(
                scenes=Timeline.request_from_payload({"scenes": scenes}).scenes,
                template=resolved.template,
                platform=resolved.platform,
                output_name=resolved.output_name,
                voiceover=resolved.voiceover,
                background_music=resolved.background_music,
                music_volume=resolved.music_volume,
                music_fade_in=resolved.music_fade_in,
                music_fade_out=resolved.music_fade_out,
                music_loop=resolved.music_loop,
                music_start=resolved.music_start,
                music_ducking=resolved.music_ducking,
                voice_volume=resolved.voice_volume,
                creative_style=resolved.creative_style,
                metadata={"creative_style": style.key, "creative_style_label": style.label},
            )
            return self._render_resolved(render_request, job_id, workdir)

    def inspect_asset(self, path: str | Path) -> dict:
        return self.assets.probe(path).__dict__

    def transcribe(
        self,
        audio_path: str | Path,
        *,
        model_size: str | None = None,
        language: str | None = None,
        task: str = "transcribe",
        beam_size: int = 5,
        vad_filter: bool = True,
        word_timestamps: bool = True,
    ) -> dict:
        return self.transcriber.transcribe(
            audio_path,
            model_size=model_size,
            language=language,
            task=task,
            beam_size=beam_size,
            vad_filter=vad_filter,
            word_timestamps=word_timestamps,
        ).as_dict()

    def _render_resolved(self, request: RenderRequest, job_id: str, workdir: Path) -> RenderResult:
        self._commands = []
        self._logs = []
        started = time.time()
        timeline = Timeline(request.scenes)
        prepared = self._prepare_scene_media(request, workdir)
        subtitles = self._write_subtitles(request, timeline, workdir)
        video = self._compose_video(prepared, request, subtitles, workdir)
        mixed = self.audio.mix(
            video,
            request.background_music,
            request.voiceover,
            workdir / "mixed.mp4",
            normalize=request.audio_normalize,
            duration=timeline.total_duration,
            music_volume=request.music_volume,
            music_fade_in=request.music_fade_in,
            music_fade_out=request.music_fade_out,
            music_loop=request.music_loop,
            music_start=request.music_start,
            music_ducking=request.music_ducking,
            voice_volume=request.voice_volume,
        )
        mixed = self._apply_watermark(mixed, request, workdir)
        optimized = self._optimize_for_platform(mixed, request, workdir)
        output = self.exports.save(optimized, job_id, request.output_name)
        profile = self._profile(request)
        metrics = {
            "render_time_seconds": round(time.time() - started, 3),
            "output_size_bytes": output.stat().st_size,
            "scene_count": len(request.scenes),
            "creative_style": request.creative_style or request.metadata.get("creative_style"),
            "scene_effects": [scene.effect for scene in request.scenes if scene.effect],
            "platform": profile.metadata(),
        }
        return RenderResult(output_path=output, commands=list(self._commands), metrics=metrics, logs=list(self._logs))

    def _prepare_scene_media(self, request: RenderRequest, workdir: Path) -> list[Path]:
        prepared: list[Path] = []
        profile = self._profile(request)
        for idx, scene in enumerate(request.scenes):
            source = Path(scene.media)
            metadata = self.assets.probe(source)
            target = workdir / f"scene_{idx:03d}.mp4"
            mime_type = metadata.mime_type or mimetypes.guess_type(str(source))[0] or ""
            if mime_type.startswith("image/") or source.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp", ".gif"}:
                self.normalizer.image_to_video(source, target, scene.duration, profile=profile)
            elif request.normalize and self.normalizer.needs_normalization(metadata, profile=profile):
                self.normalizer.normalize(source, target, scene.duration, profile=profile)
            else:
                shutil.copy2(source, target)
            if request.preview:
                preview = workdir / f"scene_{idx:03d}_preview.mp4"
                self._scale_preview(target, preview)
                target = preview
            effected = self._apply_scene_effect(target, scene.effect, workdir / f"scene_{idx:03d}_effect.mp4")
            if effected != target:
                target = effected
            prepared.append(target)
        return prepared

    def _write_subtitles(self, request: RenderRequest, timeline: Timeline, workdir: Path) -> Path | None:
        if request.auto_subtitles and request.voiceover:
            transcript = self.transcriber.transcribe(
                request.voiceover,
                model_size=request.whisper_model_size,
                language=request.subtitle_language,
                word_timestamps=True,
            )
            events = transcript.subtitle_events(prefer_words=True)
        else:
            events = self.subtitles.from_scenes(request.scenes, timeline.total_duration)
        if not events:
            return None
        if request.subtitle_format == "srt":
            return self.subtitles.write_srt(events, workdir / "captions.srt")
        return self.subtitles.write_ass(events, workdir / "captions.ass", request.template)

    def _compose_video(self, scenes: list[Path], request: RenderRequest, subtitle_path: Path | None, workdir: Path) -> Path:
        if len(scenes) == 1:
            composed = workdir / "composed.mp4"
            shutil.copy2(scenes[0], composed)
        else:
            composed = workdir / "composed.mp4"
            durations = [scene.duration for scene in request.scenes]
            transitions = [scene.transition for scene in request.scenes]
            filter_graph, final_stream = self.transitions.xfade_chain(len(scenes), durations, transitions)
            cmd = FFmpegCommand().add("-hide_banner")
            for scene in scenes:
                cmd.input(scene)
            cmd.add("-filter_complex", filter_graph)
            cmd.add("-map", final_stream, "-an", "-c:v", "libx264", "-preset", self.settings.preset, "-crf", self.settings.crf)
            command = cmd.overwrite().add(composed).build()
            self._run(command)
        if subtitle_path:
            subtitled = workdir / "subtitled.mp4"
            escaped = _ffmpeg_subtitle_path(subtitle_path)
            command = (
                FFmpegCommand()
                .add("-hide_banner")
                .input(composed)
                .add("-vf", f"subtitles='{escaped}'", "-c:a", "copy", "-c:v", "libx264", "-preset", self.settings.preset, "-crf", self.settings.crf)
                .overwrite()
                .add(subtitled)
                .build()
            )
            self._run(command)
            return subtitled
        return composed

    def _apply_watermark(self, video: Path, request: RenderRequest, workdir: Path) -> Path:
        if not request.watermark:
            return video
        watermark = Path(request.watermark)
        if not watermark.exists():
            self._logs.append(f"Watermark skipped; file not found: {watermark}")
            return video
        output = workdir / "watermarked.mp4"
        position = {
            "top-left": "20:20",
            "top-right": "W-w-20:20",
            "bottom-left": "20:H-h-20",
            "bottom-right": "W-w-20:H-h-20",
            "center": "(W-w)/2:(H-h)/2",
        }.get(request.watermark_position, "W-w-20:H-h-20")
        command = (
            FFmpegCommand()
            .add("-hide_banner")
            .input(video)
            .input(watermark)
            .add("-filter_complex", f"[1:v]scale=iw*0.22:-1[wm];[0:v][wm]overlay={position}")
            .add("-c:a", "copy", "-c:v", "libx264", "-preset", self.settings.preset, "-crf", self.settings.crf)
            .overwrite()
            .add(output)
            .build()
        )
        self._run(command)
        return output

    def _optimize_for_platform(self, video: Path, request: RenderRequest, workdir: Path) -> Path:
        profile = self._profile(request)
        output = workdir / "platform_optimized.mp4"
        vf = (
            f"scale={profile.width}:{profile.height}:force_original_aspect_ratio=increase,"
            f"crop={profile.width}:{profile.height},fps={profile.fps},format=yuv420p"
        )
        command_builder = (
            FFmpegCommand()
            .add("-hide_banner")
            .input(video)
            .add(
                "-vf",
                vf,
                "-c:v",
                profile.video_codec,
                "-profile:v",
                "high",
                "-pix_fmt",
                "yuv420p",
                "-preset",
                self.settings.preset,
                "-crf",
                profile.crf,
                "-r",
                profile.fps,
                "-g",
                max(1, profile.fps * 2),
            )
        )
        if profile.maxrate:
            command_builder.add("-maxrate", profile.maxrate)
        if profile.bufsize:
            command_builder.add("-bufsize", profile.bufsize)
        command = (
            command_builder.add(
                "-c:a",
                profile.audio_codec,
                "-b:a",
                profile.audio_bitrate,
                "-ar",
                profile.audio_sample_rate,
                "-ac",
                "2",
                "-movflags",
                "+faststart",
                "-shortest",
            )
            .overwrite()
            .add(output)
            .build()
        )
        self._run(command)
        return output

    def _profile(self, request: RenderRequest) -> PlatformProfile:
        profile_key = request.platform or request.metadata.get("platform") or request.metadata.get("target_platform")
        return get_platform_profile(str(profile_key) if profile_key else None)


    def _scale_preview(self, source: Path, output: Path) -> None:
        command = (
            FFmpegCommand()
            .add("-hide_banner")
            .input(source)
            .add("-vf", "scale=540:960:force_original_aspect_ratio=decrease,pad=540:960:(ow-iw)/2:(oh-ih)/2")
            .add("-c:v", "libx264", "-preset", "ultrafast", "-crf", "30", "-c:a", "aac")
            .overwrite()
            .add(output)
            .build()
        )
        self._run(command)

    def _apply_scene_effect(self, source: Path, effect: str | None, output: Path) -> Path:
        vf = scene_effect_filter(effect)
        if not vf:
            return source
        command = (
            FFmpegCommand()
            .add("-hide_banner")
            .input(source)
            .add("-vf", vf, "-c:v", "libx264", "-preset", self.settings.preset, "-crf", self.settings.crf, "-c:a", "copy")
            .overwrite()
            .add(output)
            .build()
        )
        self._run(command)
        self._logs.append(f"Applied scene effect '{effect}' to {source.name}")
        return output

    def _run(self, command: list[str]) -> None:
        result = self.runner.run(command)
        if result.stderr:
            self._logs.append(result.stderr[-4000:])

    def _record_command(self, command: list[str]) -> None:
        self._commands.append(command)


def _split_script(script: str, chunks: int) -> list[str]:
    words = script.split()
    if chunks <= 0:
        return []
    size = max(1, round(len(words) / chunks))
    return [" ".join(words[i : i + size]) for i in range(0, len(words), size)][:chunks]


def _ffmpeg_subtitle_path(path: Path) -> str:
    return str(path).replace("\\", "/").replace(":", r"\:")