"""Wan2.1 1.3B text-to-video starter (diffusers).

Bezpečnost:
- torch.load vždy s weights_only=True
- Žádný trust_remote_code=True
- Revize modelu pinována na konkrétní SHA
- SHA256 ověření stažených souborů
- GPU vyžadováno (~8 GB VRAM), CPU není praktické
"""

from __future__ import annotations

import hashlib
import os
import sys
from pathlib import Path

import torch
from diffusers import WanPipeline  # typový placeholder, případně Wan2.1Pipeline

HF_TOKEN = os.environ.get("HF_TOKEN")  # výhradně z env, žádné výchozí hodnoty
MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B"  # Apache 2.0
REVISION = "<revision_sha>"  # nutné pinovat na konkrétní commit SHA
MODEL_FILE = "diffusion_pytorch_model.safetensors"
EXPECTED_SHA256 = "<sha256_digest>"  # doplnit po stažení
MODELS_DIR = Path(os.environ.get("DEV_ENV_MODELS_DIR", "models"))


def _verify_sha256(filepath: Path, expected: str) -> str:
    sha = hashlib.sha256()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            sha.update(chunk)
    digest = sha.hexdigest()
    if digest != expected:
        print(f"CHYBA: SHA256 neshoda pro {filepath.name}: {digest} != {expected}", file=sys.stderr)
        sys.exit(1)
    return digest


def _download_weight(repo_id: str, filename: str, revision: str, expected_hash: str) -> Path:
    from huggingface_hub import hf_hub_download

    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    downloaded = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        revision=revision,
        token=HF_TOKEN,
        local_dir=str(MODELS_DIR),
        local_dir_use_symlinks=False,
        trust_remote_code=False,  # NIKDY True
    )
    actual = _verify_sha256(Path(downloaded), expected_hash)
    print(f"SHA256 ověřeno: {actual}")
    return Path(downloaded)


def generate_video(prompt: str, output_path: str = "output.mp4") -> None:
    if not torch.cuda.is_available():
        print("VAROVÁNÍ: GPU není k dispozici — Wan2.1 vyžaduje ~8 GB VRAM. CPU běh není praktický.", file=sys.stderr)

    weight_path = _download_weight(MODEL_ID, MODEL_FILE, REVISION, EXPECTED_SHA256)

    # Načtení modelu s weights_only=True (žádné pickle exécutující kód)
    state_dict = torch.load(weight_path, map_location="cpu", weights_only=True)

    pipeline = WanPipeline.from_pretrained(
        MODEL_ID,
        revision=REVISION,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        token=HF_TOKEN,
        trust_remote_code=False,
    )
    pipeline.to("cuda" if torch.cuda.is_available() else "cpu")

    result = pipeline(prompt=prompt)
    result.videos[0].save(output_path)
    print(f"Video uloženo: {output_path}")


if __name__ == "__main__":
    generate_video("A cinematic video of a cat walking in a garden.")