from __future__ import annotations

import os
import subprocess
import tempfile
import threading
from pathlib import Path
from typing import Any


_MODEL: Any | None = None
_MODEL_LOCK = threading.Lock()


def _cache_dir() -> Path:
    preferred = Path(os.getenv("TRIBE_CACHE_DIR", "/data/clawbrain-cache"))
    if preferred.parent.exists() and os.access(preferred.parent, os.W_OK):
        preferred.mkdir(parents=True, exist_ok=True)
        return preferred
    fallback = Path(".cache/clawbrain")
    fallback.mkdir(parents=True, exist_ok=True)
    return fallback


def get_model() -> Any:
    global _MODEL
    if _MODEL is not None:
        return _MODEL
    with _MODEL_LOCK:
        if _MODEL is None:
            try:
                from tribev2.demo_utils import TribeModel
            except ImportError as exc:
                raise RuntimeError(
                    "TRIBE v2 is not installed. Install the Space requirements first."
                ) from exc
            _MODEL = TribeModel.from_pretrained(
                "facebook/tribev2",
                cache_folder=_cache_dir(),
                device=os.getenv("TRIBE_DEVICE", "auto"),
                config_update={
                    # ZeroGPU executes decorated functions in a daemon worker;
                    # PyTorch multiprocessing workers cannot be spawned there.
                    "data.num_workers": 0,
                    # One short clip at a time keeps host and VRAM peaks stable.
                    "data.batch_size": 1,
                },
            )
    return _MODEL


def predict_video(video_path: str, progress: Any | None = None) -> tuple[Any, Any]:
    if progress:
        progress(0.04, desc="Preparing gameplay video")

    # TribeModel.get_events_dataframe() always launches WhisperX through `uvx`
    # to add word-level events. For gameplay captures, the visual and audio
    # streams are the primary signal; skipping transcription avoids a second
    # multi-GB ASR stack and keeps the Space independent of gated Llama access.
    from tribev2.demo_utils import get_audio_and_text_events
    import pandas as pd

    # Mobile captures commonly arrive at 30-60 FPS, while the brain encoder's
    # temporal signal does not require every display frame. Downsampling to
    # 8 FPS preserves the complete scene and audio timeline while reducing
    # V-JEPA work enough to fit a shared ZeroGPU reservation.
    with tempfile.TemporaryDirectory(prefix="clawbrain-video-") as tmp_dir:
        prepared_video = Path(tmp_dir) / "gameplay-8fps.mp4"
        command = [
            "ffmpeg",
            "-y",
            "-i",
            str(Path(video_path).resolve()),
            "-vf",
            "fps=8",
            "-c:v",
            "libx264",
            "-preset",
            "veryfast",
            "-crf",
            "20",
            "-pix_fmt",
            "yuv420p",
            "-c:a",
            "aac",
            "-b:a",
            "128k",
            str(prepared_video),
        ]
        result = subprocess.run(command, capture_output=True, text=True)
        if result.returncode != 0:
            raise RuntimeError(f"Video preparation failed: {result.stderr[-600:]}")

        if progress:
            progress(0.08, desc="Loading TRIBE v2")
        model = get_model()
        if progress:
            progress(0.20, desc="Extracting video and audio events")
        events = get_audio_and_text_events(
            pd.DataFrame(
                [
                    {
                        "type": "Video",
                        "filepath": str(prepared_video),
                        "start": 0,
                        "timeline": "default",
                        "subject": "default",
                    }
                ]
            ),
            audio_only=True,
        )
        if progress:
            progress(0.42, desc="Encoding multimodal content")
        predictions, segments = model.predict(events=events)
        if progress:
            progress(0.90, desc="Building cortical response timeline")
        return predictions, segments