# hd_specialist.py
#
# Copyright (C) 2025 Carlos Rodrigues dos Santos
#
# This file implements the HD Specialist (Δ+), which uses the SeedVR model
# for video super-resolution. It's designed to be called by the ADUC orchestrator
# to perform the final HD mastering pass on a generated video. It manages the
# loading/unloading of the heavy SeedVR models to conserve VRAM and can switch
# between different model sizes (e.g., 3B and 7B).

import torch
import imageio
import os
import gc
import logging
import numpy as np
from PIL import Image
from tqdm import tqdm
import shlex
import subprocess
from pathlib import Path
from urllib.parse import urlparse
from torch.hub import download_url_to_file
from omegaconf import OmegaConf
import mediapy
from einops import rearrange

# Assuming these files are in the project structure
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 projects.video_diffusion_sr.color_fix import wavelet_reconstruction
from torchvision.transforms import Compose, Lambda, Normalize
from torchvision.io.video import read_video

logger = logging.getLogger(__name__)

def _load_file_from_url(url, model_dir='./', file_name=None):
    """Helper function to download files from a URL to a local directory."""
    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 HDSpecialist:
    """
    Implements the HD Specialist (Δ+) using the SeedVR infrastructure.
    Manages model loading, inference, and memory on demand.
    """
    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
        logger.info("HD Specialist (SeedVR) initialized. Model will be loaded on demand.")

    def _download_models(self):
        """Downloads the necessary checkpoints for SeedVR2."""
        logger.info("Verifying and downloading SeedVR2 models...")
        ckpt_dir = Path('./ckpts')
        ckpt_dir.mkdir(exist_ok=True)

        pretrain_model_urls = {
            'vae': '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='./ckpts/')
            
        logger.info("SeedVR2 models downloaded successfully.")

    def _initialize_runner(self, model_version: str):
        """Loads and configures the SeedVR model on demand based on the selected version."""
        if self.runner is not None:
            return

        self._download_models()

        logger.info(f"Initializing SeedVR2 {model_version} runner...")
        if model_version == '3B':
            config_path = os.path.join('./configs_3b', 'main.yaml')
            checkpoint_path = './ckpts/seedvr2_ema_3b.pth'
        elif model_version == '7B':
            config_path = os.path.join('./configs_7b', 'main.yaml')
            checkpoint_path = './ckpts/seedvr2_ema_7b.pth'
        else:
            raise ValueError(f"Unsupported SeedVR model version: {model_version}")

        config = load_config(config_path)
        
        self.runner = VideoDiffusionInfer(config)
        OmegaConf.set_readonly(self.runner.config, False)
        
        self.runner.configure_dit_model(device=self.device, checkpoint=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):
        """Removes the runner from VRAM to free resources."""
        if self.runner is not None:
            del self.runner
            self.runner = None
            gc.collect()
            torch.cuda.empty_cache()
            self.is_initialized = False
            logger.info("SeedVR2 runner unloaded from VRAM.")

    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)

            # --- Adapted inference logic from SeedVR scripts ---
            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)

            text_pos_embeds = torch.load('./ckpts/pos_emb.pt').to(self.device)
            text_neg_embeds = torch.load('./ckpts/neg_emb.pt').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]: # if generated frames are less
                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
hd_specialist_singleton = HDSpecialist()