# audio_specialist.py (Versão final para áudio dinâmico por fragmento)
# Especialista ADUC para geração de áudio, com gerenciamento de memória GPU.

import torch
import logging
import subprocess
import os
import time
import yaml
import gc
from pathlib import Path
import gradio as gr

# Importa as classes e funções necessárias do MMAudio
try:
    from mmaudio.eval_utils import ModelConfig, all_model_cfg, generate as mmaudio_generate, load_video, make_video
    from mmaudio.model.flow_matching import FlowMatching
    from mmaudio.model.networks import MMAudio, get_my_mmaudio
    from mmaudio.model.utils.features_utils import FeaturesUtils
    from mmaudio.model.sequence_config import SequenceConfig
except ImportError:
    raise ImportError("MMAudio não foi encontrado. Por favor, instale-o a partir do GitHub: git+https://github.com/hkchengrex/MMAudio.git")

logger = logging.getLogger(__name__)

class AudioSpecialist:
    def __init__(self, workspace_dir):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.cpu_device = torch.device("cpu")
        self.dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
        self.workspace_dir = workspace_dir
        
        self.model_config: ModelConfig = all_model_cfg['large_44k_v2']
        self.net: MMAudio = None
        self.feature_utils: FeaturesUtils = None
        self.seq_cfg: SequenceConfig = None
        
        self._load_models_to_cpu()

    def _load_models_to_cpu(self):
        try:
            logger.info("Verificando e baixando modelos MMAudio, se necessário...")
            self.model_config.download_if_needed()
            
            self.seq_cfg = self.model_config.seq_cfg
            
            logger.info(f"Carregando modelo MMAudio: {self.model_config.model_name} para a CPU...")
            self.net = get_my_mmaudio(self.model_config.model_name).eval()
            self.net.load_weights(torch.load(self.model_config.model_path, map_location=self.cpu_device, weights_only=True))
            
            logger.info("Carregando utilitários de features do MMAudio para a CPU...")
            self.feature_utils = FeaturesUtils(
                tod_vae_ckpt=self.model_config.vae_path,
                synchformer_ckpt=self.model_config.synchformer_ckpt,
                enable_conditions=True,
                mode=self.model_config.mode,
                bigvgan_vocoder_ckpt=self.model_config.bigvgan_16k_path,
                need_vae_encoder=False
            )
            self.feature_utils = self.feature_utils.eval()
            self.net.to(self.cpu_device)
            self.feature_utils.to(self.cpu_device)
            logger.info("Especialista de áudio pronto na CPU.")
        except Exception as e:
            logger.error(f"Falha ao carregar modelos de áudio: {e}", exc_info=True)
            self.net = None

    def to_gpu(self):
        if self.device == 'cpu': return
        logger.info(f"Movendo especialista de áudio para a GPU ({self.device})...")
        self.net.to(self.device, self.dtype)
        self.feature_utils.to(self.device, self.dtype)

    def to_cpu(self):
        if self.device == 'cpu': return
        logger.info("Descarregando especialista de áudio da GPU...")
        self.net.to(self.cpu_device)
        self.feature_utils.to(self.cpu_device)
        gc.collect()
        if torch.cuda.is_available(): torch.cuda.empty_cache()

    def generate_audio_for_video(self, video_path: str, prompt: str, negative_prompt: str, duration_seconds: float) -> str:
        if self.net is None:
            raise gr.Error("Modelo MMAudio não está carregado. Não é possível gerar áudio.")

        logger.info("------------------------------------------------------")
        logger.info("--- Gerando Áudio para Fragmento de Vídeo ---")
        logger.info(f"--- Vídeo Fragmento: {os.path.basename(video_path)}")
        logger.info(f"--- Duração: {duration_seconds:.2f}s")
        logger.info(f"--- Prompt (Descrição da Cena): '{prompt}'")
        
        if duration_seconds < 1:
            logger.warning("Fragmento muito curto (<1s). Retornando vídeo silencioso.")
            logger.info("------------------------------------------------------")
            return video_path

        if self.device == 'cpu':
            logger.warning("Gerando áudio na CPU. Isso pode ser muito lento.")

        try:
            self.to_gpu()
            with torch.no_grad():
                rng = torch.Generator(device=self.device).manual_seed(int(time.time()))
                fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=25)
                
                video_info = load_video(Path(video_path), duration_seconds)
                self.seq_cfg.duration = video_info.duration_sec
                self.net.update_seq_lengths(self.seq_cfg.latent_seq_len, self.seq_cfg.clip_seq_len, self.seq_cfg.sync_seq_len)
                
                audios = mmaudio_generate(
                    clip_video=video_info.clip_frames.unsqueeze(0),
                    sync_video=video_info.sync_frames.unsqueeze(0),
                    text=[prompt],
                    negative_text=[negative_prompt],
                    feature_utils=self.feature_utils,
                    net=self.net,
                    fm=fm,
                    rng=rng,
                    cfg_strength=4.5
                )
                audio_waveform = audios.float().cpu()[0]
                
                fragment_name = Path(video_path).stem
                output_video_path = os.path.join(self.workspace_dir, f"{fragment_name}_com_audio.mp4")
                
                make_video(video_info, Path(output_video_path), audio_waveform, sampling_rate=self.seq_cfg.sampling_rate)
                logger.info(f"--- Fragmento com áudio salvo em: {os.path.basename(output_video_path)}")
                logger.info("------------------------------------------------------")
                return output_video_path
        finally:
            self.to_cpu()

# Singleton instantiation
try:
    with open("config.yaml", 'r') as f:
        config = yaml.safe_load(f)
    WORKSPACE_DIR = config['application']['workspace_dir']
    audio_specialist_singleton = AudioSpecialist(workspace_dir=WORKSPACE_DIR)
except Exception as e:
    logger.error(f"Não foi possível inicializar o AudioSpecialist: {e}")
    audio_specialist_singleton = None