"""Frame export (last-frame extraction from a generated video).

`export_last_frame` is bound to the "Export Last Frame" button in the UI.
It validates that the input path lies inside the managed task temp root or the
legacy `/tmp/h3-outputs` directory, then runs ffmpeg to extract the final frame.

This is the only place that exposes the disk to client-supplied paths; the
path validation is intentionally strict.
"""
from __future__ import annotations

import os
import subprocess
import tempfile
import uuid

import gradio as gr

from task import _TASK_TEMP_ROOT


def export_last_frame(video_path):
    """Extract the last frame from a generated video and return it as an image file.

    Path is validated to lie under the managed task/temp root; arbitrary client paths are rejected.
    Destination is unique per call.
    """
    if not video_path:
        raise gr.Error("No video to extract from.")

    # Resolve and constrain to managed roots (task temp or legacy h3-outputs).
    try:
        resolved = os.path.realpath(video_path)
    except OSError:
        raise gr.Error("Invalid video path.")
    # Only managed task/output roots — never the bare system tempdir (symlink surface).
    allowed_roots = [
        os.path.realpath(_TASK_TEMP_ROOT),
        os.path.realpath(os.path.join(tempfile.gettempdir(), "h3-outputs")),
    ]
    if not any(resolved == root or resolved.startswith(root + os.sep) for root in allowed_roots):
        raise gr.Error("Video path is outside the managed output area.")
    # Unique destination under the same parent or task intermediates.
    parent = os.path.dirname(resolved)
    out = os.path.join(parent, f"lastframe-{uuid.uuid4().hex[:12]}.png")
    result = subprocess.run(
        ["ffmpeg", "-sseof", "-0.1", "-i", resolved, "-frames:v", "1", "-y", out],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise gr.Error(f"ffmpeg failed: {result.stderr[-300:]}")
    return out


# ── Durable video finalization (display + download) ─────────────────────────

# Public-ish mirror of task outputs. Gradio serves these more reliably than deep
# task dirs, and clients can still download after the request finishes.
_H3_OUTPUT_ROOT = os.environ.get(
    "H3_OUTPUT_ROOT", os.path.join(tempfile.gettempdir(), "h3-outputs")
)


def _ensure_output_roots():
    for d in (_H3_OUTPUT_ROOT, _TASK_TEMP_ROOT):
        try:
            os.makedirs(d, exist_ok=True)
        except OSError:
            pass


def _finalize_output_video(frames, audio, sampling_rate, *, task=None, fps=24):
    """Mux frames (+ optional audio) to a durable mp4 and return an absolute path.

    Hardens Gradio Video / Download against:
    - missing/empty files after encode
    - encode failures with audio (retries silent)
    - non-absolute paths Gradio cannot serve
    - deep task-only paths that disappear from the UI file server
    """
    import shutil
    from diffusers.utils import encode_video as _encode_video

    _ensure_output_roots()

    # Normalize frames to a numpy-friendly object encode_video accepts.
    if hasattr(frames, "detach"):
        frames = frames.detach().cpu().numpy()
    if frames is None:
        raise gr.Error("No frames were produced — nothing to encode.")

    tid = (task.task_id if task is not None else None) or uuid.uuid4().hex
    work = (task.work_dir if task is not None and getattr(task, "work_dir", None) else None)
    if work:
        try:
            os.makedirs(work, exist_ok=True)
        except OSError:
            work = None
    primary = os.path.join(work or _H3_OUTPUT_ROOT, f"h3-{tid}.mp4")
    mirror = os.path.join(_H3_OUTPUT_ROOT, f"h3-{tid}.mp4")

    def _try_encode(path, with_audio: bool):
        kwargs = dict(fps=fps, output_path=path)
        if with_audio and audio is not None:
            kwargs["audio"] = audio
            kwargs["audio_sample_rate"] = sampling_rate
        _encode_video(frames, **kwargs)

    last_err = None
    encoded = False
    for with_audio in (True, False):
        try:
            _try_encode(primary, with_audio=with_audio)
            encoded = True
            if not with_audio and audio is not None:
                # Silent fallback — still deliver a playable file.
                pass
            break
        except Exception as e:
            last_err = e
            try:
                if os.path.isfile(primary):
                    os.remove(primary)
            except OSError:
                pass
            continue
    if not encoded:
        raise gr.Error(
            f"Failed to mux the video file ({type(last_err).__name__}: {last_err}). "
            "Try a shorter duration or fewer steps."
        )

    if not os.path.isfile(primary) or os.path.getsize(primary) < 1024:
        raise gr.Error(
            "Video file was not written correctly (missing or under 1 KB). "
            "Please retry the generation."
        )

    # Mirror into h3-outputs so Gradio always has a stable root to serve from.
    try:
        if os.path.realpath(primary) != os.path.realpath(mirror):
            shutil.copy2(primary, mirror)
            serve_path = mirror
        else:
            serve_path = primary
    except OSError:
        serve_path = primary

    serve_path = os.path.realpath(serve_path)
    if task is not None:
        task.output_path = serve_path
    return serve_path