"""Wrapper around the `videoseal` package.

The upstream package on PyPI (videoseal 1.0.1):
- exposes a single `videoseal.load(model_card='videoseal')` entry point.
- the loader expects `configs/attenuation.yaml` to exist relative to the CWD
  (the upstream code path was written for a repo checkout). We bundle a copy
  under `setup/src/configs/` and ensure it's present in CWD on import.
- `model.embed(imgs, msgs=..., is_video=True)` -> {"imgs_w": F×3×H×W, "msgs": ...}
- `model.extract_message(imgs, aggregation="avg")` -> Tensor[1, K] of 0/1 bits.
- Default model card `videoseal_1.0` ships nbits=256.
"""

from __future__ import annotations

import os
import shutil
import threading
from pathlib import Path
from typing import Optional

import numpy as np
import torch
import torchvision.io

from . import config

_HERE = Path(__file__).resolve().parent
_BUNDLED_CONFIGS = _HERE / "configs"


def _ensure_runtime_layout() -> None:
    """videoseal.load() reads two paths relative to the CWD:

      - `videoseal/cards/<model_card>.yaml`  (Path("videoseal/cards") hardcoded)
      - `configs/attenuation.yaml`           (read from the model card YAML)

    Both happen because the upstream code was written for a repo checkout.
    We mirror the bundled `configs/attenuation.yaml` and the model cards from
    the installed package into the CWD on demand so the loader resolves them."""
    target_attn = Path("configs") / "attenuation.yaml"
    if not target_attn.exists():
        target_attn.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(_BUNDLED_CONFIGS / "attenuation.yaml", target_attn)

    target_cards = Path("videoseal") / "cards"
    if not target_cards.exists():
        import videoseal as _vs
        src_cards = Path(_vs.__file__).resolve().parent / "cards"
        target_cards.mkdir(parents=True, exist_ok=True)
        for card in src_cards.glob("*.yaml"):
            shutil.copy(card, target_cards / card.name)


_MODEL = None
_MODEL_LOCK = threading.Lock()


def get_model():
    """Return a cached videoseal model. Loads weights from the upstream URL
    on first call (cached to `./ckpts/` by the library itself)."""
    global _MODEL
    if _MODEL is not None:
        return _MODEL
    with _MODEL_LOCK:
        if _MODEL is None:
            import videoseal
            _ensure_runtime_layout()
            model = videoseal.load("videoseal")
            model.eval()
            _MODEL = model
    return _MODEL


def _read_video_frames(path: Path, max_frames: Optional[int] = None,
                       subsample: int = 1, target_side: Optional[int] = None) -> torch.Tensor:
    """Read MP4/MOV/WebM into a (T, 3, H, W) float tensor in [0, 1]."""
    video, _, _ = torchvision.io.read_video(str(path), pts_unit="sec")
    if video.numel() == 0:
        raise ValueError("video has zero frames")
    video = video.permute(0, 3, 1, 2).float() / 255.0
    if subsample > 1:
        video = video[::subsample]
    if max_frames is not None and video.shape[0] > max_frames:
        video = video[:max_frames]
    if target_side is not None:
        _, _, h, w = video.shape
        scale = target_side / max(h, w)
        new_h = max(1, int(round(h * scale)))
        new_w = max(1, int(round(w * scale)))
        if (new_h, new_w) != (h, w):
            video = torch.nn.functional.interpolate(
                video, size=(new_h, new_w), mode="bilinear", align_corners=False
            )
    return video


def embed_into_video(in_path: Path, out_path: Path, payload_bits: np.ndarray,
                     fps: int = 24) -> None:
    """Embed a 256-bit payload into the video at `in_path`, write to `out_path`."""
    frames = _read_video_frames(in_path, target_side=None)
    msg = torch.tensor(payload_bits, dtype=torch.float32).unsqueeze(0)
    outputs = get_model().embed(frames, msgs=msg, is_video=True)
    watermarked = outputs["imgs_w"].clamp(0, 1)
    watermarked = (watermarked * 255).to(torch.uint8).permute(0, 2, 3, 1).cpu()
    torchvision.io.write_video(
        str(out_path),
        watermarked,
        fps=fps,
        video_codec="libx264",
        options={"crf": "10", "preset": "medium"},
    )


def extract_payload(path: Path) -> np.ndarray:
    """Run videoseal extraction on the video at `path`. Returns a (256,) uint8 bit array."""
    frames = _read_video_frames(
        path,
        subsample=config.FRAME_SUBSAMPLE,
        target_side=config.EXTRACT_SIDE,
        max_frames=64,
    )
    with torch.no_grad():
        msg = get_model().extract_message(frames, aggregation="avg")
    bits = msg.detach().cpu().to(torch.uint8).squeeze().numpy().reshape(-1)
    if bits.size < 256:
        raise RuntimeError(f"videoseal returned {bits.size} bits, expected >= 256")
    return bits[:256]
