# managers/seedvr_manager.py
#
# Copyright (C) 2025 Carlos Rodrigues dos Santos
#
# Version: 2.3.3
#
# This version adds a monkey patch to disable torch.distributed.barrier calls
# within the SeedVR library, allowing it to run in a single-GPU inference mode
# without raising a "process group not initialized" error.

import torch
import torch.distributed as dist
import os
import gc
import logging
import sys
import subprocess
from pathlib import Path
from urllib.parse import urlparse
from torch.hub import download_url_to_file
import gradio as gr
import mediapy
from einops import rearrange

from tools.tensor_utils import wavelet_reconstruction

logger = logging.getLogger(__name__)

# --- Dependency Management ---
DEPS_DIR = Path("./deps")
SEEDVR_REPO_DIR = DEPS_DIR / "SeedVR"
SEEDVR_REPO_URL = "https://github.com/ByteDance-Seed/SeedVR.git"
VAE_CONFIG_URL = "https://raw.githubusercontent.com/ByteDance-Seed/SeedVR/main/models/video_vae_v3/s8_c16_t4_inflation_sd3.yaml"

def setup_seedvr_dependencies():
    """Ensures the SeedVR repository is cloned and available in the sys.path."""
    if not SEEDVR_REPO_DIR.exists():
        logger.info(f"SeedVR repository not found at '{SEEDVR_REPO_DIR}'. Cloning from GitHub...")
        try:
            DEPS_DIR.mkdir(exist_ok=True)
            subprocess.run(["git", "clone", "--depth", "1", SEEDVR_REPO_URL, str(SEEDVR_REPO_DIR)], check=True, capture_output=True, text=True)
            logger.info("SeedVR repository cloned successfully.")
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to clone SeedVR repository. Git stderr: {e.stderr}")
            raise RuntimeError("Could not clone the required SeedVR dependency from GitHub.")
    else:
        logger.info("Found local SeedVR repository.")
    
    if str(SEEDVR_REPO_DIR.resolve()) not in sys.path:
        sys.path.insert(0, str(SEEDVR_REPO_DIR.resolve()))
        logger.info(f"Added '{SEEDVR_REPO_DIR.resolve()}' to sys.path.")

setup_seedvr_dependencies()

from projects.video_diffusion_sr.infer import VideoDiffusionInfer
from common.config import load_config
from common.seed import set_seed
from data.image.transforms.divisible_crop import DivisibleCrop
from data.image.transforms.na_resize import NaResize
from data.video.transforms.rearrange import Rearrange
from torchvision.transforms import Compose, Lambda, Normalize
from torchvision.io.video import read_video
from omegaconf import OmegaConf


def _load_file_from_url(url, model_dir='./', file_name=None):
    os.makedirs(model_dir, exist_ok=True)
    filename = file_name or os.path.basename(urlparse(url).path)
    cached_file = os.path.abspath(os.path.join(model_dir, filename))
    if not os.path.exists(cached_file):
        logger.info(f'Downloading: "{url}" to {cached_file}')
        download_url_to_file(url, cached_file, hash_prefix=None, progress=True)
    return cached_file

class SeedVrManager:
    """Manages the SeedVR model for HD Mastering tasks."""
    def __init__(self, workspace_dir="deformes_workspace"):
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.runner = None
        self.workspace_dir = workspace_dir
        self.is_initialized = False
        self._original_barrier = None # To store the original distributed barrier function
        logger.info("SeedVrManager initialized. Model will be loaded on demand.")

    def _download_models_and_configs(self):
        """Downloads the necessary checkpoints AND the missing VAE config file."""
        logger.info("Verifying and downloading SeedVR2 models and configs...")
        ckpt_dir = SEEDVR_REPO_DIR / 'ckpts'
        config_dir = SEEDVR_REPO_DIR / 'configs' / 'vae'
        ckpt_dir.mkdir(exist_ok=True)
        config_dir.mkdir(parents=True, exist_ok=True)
        _load_file_from_url(url=VAE_CONFIG_URL, model_dir=str(config_dir))
        pretrain_model_urls = {
            'vae_ckpt': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/ema_vae.pth',
            'dit_3b': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/seedvr2_ema_3b.pth',
            'dit_7b': 'https://huggingface.co/ByteDance-Seed/SeedVR2-7B/resolve/main/seedvr2_ema_7b.pth',
            'pos_emb': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/pos_emb.pt',
            'neg_emb': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/neg_emb.pt'
        }
        for key, url in pretrain_model_urls.items():
            _load_file_from_url(url=url, model_dir=str(ckpt_dir))
        logger.info("SeedVR2 models and configs downloaded successfully.")

    def _initialize_runner(self, model_version: str):
        """Loads and configures the SeedVR model, with patches for single-GPU inference."""
        if self.runner is not None: return
        self._download_models_and_configs()

        if dist.is_available() and not dist.is_initialized():
            logger.info("Applying patch to disable torch.distributed.barrier for single-GPU inference.")
            self._original_barrier = dist.barrier
            dist.barrier = lambda *args, **kwargs: None
        
        logger.info(f"Initializing SeedVR2 {model_version} runner...")
        if model_version == '3B':
            config_path = SEEDVR_REPO_DIR / 'configs_3b' / 'main.yaml'
            checkpoint_path = SEEDVR_REPO_DIR / 'ckpts' / 'seedvr2_ema_3b.pth'
        elif model_version == '7B':
            config_path = SEEDVR_REPO_DIR / 'configs_7b' / 'main.yaml'
            checkpoint_path = SEEDVR_REPO_DIR / 'ckpts' / 'seedvr2_ema_7b.pth'
        else:
            raise ValueError(f"Unsupported SeedVR model version: {model_version}")

        try:
            config = load_config(str(config_path))
        except FileNotFoundError:
            logger.warning("Caught expected FileNotFoundError. Loading config manually.")
            config = OmegaConf.load(str(config_path))
            correct_vae_config_path = SEEDVR_REPO_DIR / 'configs' / 'vae' / 's8_c16_t4_inflation_sd3.yaml'
            vae_config = OmegaConf.load(str(correct_vae_config_path))
            config.vae = vae_config
            logger.info("Configuration loaded and patched manually.")

        self.runner = VideoDiffusionInfer(config)
        OmegaConf.set_readonly(self.runner.config, False)
        self.runner.configure_dit_model(device=self.device, checkpoint=str(checkpoint_path))
        self.runner.configure_vae_model()
        if hasattr(self.runner.vae, "set_memory_limit"):
            self.runner.vae.set_memory_limit(**self.runner.config.vae.memory_limit)
        self.is_initialized = True
        logger.info(f"Runner for SeedVR2 {model_version} initialized and ready.")
    
    def _unload_runner(self):
        """Unloads the runner from VRAM and restores any applied patches."""
        if self.runner is not None:
            del self.runner; self.runner = None
            gc.collect(); torch.cuda.empty_cache()
            self.is_initialized = False
            logger.info("SeedVR runner unloaded from VRAM.")
        
        if self._original_barrier is not None:
            logger.info("Restoring original torch.distributed.barrier function.")
            dist.barrier = self._original_barrier
            self._original_barrier = None

    def process_video(self, input_video_path: str, output_video_path: str, prompt: str,
                      model_version: str = '3B', steps: int = 50, seed: int = 666, 
                      progress: gr.Progress = None) -> str:
        """Applies HD enhancement to a video using the SeedVR logic."""
        try:
            self._initialize_runner(model_version)
            set_seed(seed, same_across_ranks=True)
            self.runner.config.diffusion.timesteps.sampling.steps = steps
            self.runner.configure_diffusion()
            video_tensor = read_video(input_video_path, output_format="TCHW")[0] / 255.0
            res_h, res_w = video_tensor.shape[-2:]
            video_transform = Compose([
                NaResize(resolution=(res_h * res_w) ** 0.5, mode="area", downsample_only=False),
                Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
                DivisibleCrop((16, 16)),
                Normalize(0.5, 0.5),
                Rearrange("t c h w -> c t h w"),
            ])
            cond_latents = [video_transform(video_tensor.to(self.device))]
            input_videos = cond_latents
            self.runner.dit.to("cpu")
            self.runner.vae.to(self.device)
            cond_latents = self.runner.vae_encode(cond_latents)
            self.runner.vae.to("cpu"); gc.collect(); torch.cuda.empty_cache()
            self.runner.dit.to(self.device)
            pos_emb_path = SEEDVR_REPO_DIR / 'ckpts' / 'pos_emb.pt'
            neg_emb_path = SEEDVR_REPO_DIR / 'ckpts' / 'neg_emb.pt'
            text_pos_embeds = torch.load(pos_emb_path).to(self.device)
            text_neg_embeds = torch.load(neg_emb_path).to(self.device)
            text_embeds_dict = {"texts_pos": [text_pos_embeds], "texts_neg": [text_neg_embeds]}
            noises = [torch.randn_like(latent) for latent in cond_latents]
            conditions = [self.runner.get_condition(noise, latent_blur=latent, task="sr") for noise, latent in zip(noises, cond_latents)]
            with torch.no_grad(), torch.autocast("cuda", torch.bfloat16, enabled=True):
                video_tensors = self.runner.inference(noises=noises, conditions=conditions, dit_offload=True, **text_embeds_dict)
            self.runner.dit.to("cpu"); gc.collect(); torch.cuda.empty_cache()
            self.runner.vae.to(self.device)
            samples = self.runner.vae_decode(video_tensors)
            final_sample = samples[0]
            input_video_sample = input_videos[0]
            if final_sample.shape[1] < input_video_sample.shape[1]:
                input_video_sample = input_video_sample[:, :final_sample.shape[1]]
            final_sample = wavelet_reconstruction(rearrange(final_sample, "c t h w -> t c h w"), rearrange(input_video_sample, "c t h w -> t c h w"))
            final_sample = rearrange(final_sample, "t c h w -> t h w c")
            final_sample = final_sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round()
            final_sample_np = final_sample.to(torch.uint8).cpu().numpy()
            mediapy.write_video(output_video_path, final_sample_np, fps=24)
            logger.info(f"HD Mastered video saved to: {output_video_path}")
            return output_video_path
        finally:
            self._unload_runner()

# --- Singleton Instance ---
seedvr_manager_singleton = SeedVrManager()