# ltx_server_refactored.py — VideoService (Modular Version)

# --- 0. WARNINGS E AMBIENTE ---
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", message=".*")
from huggingface_hub import logging
logging.set_verbosity_error()
logging.set_verbosity_warning()
logging.set_verbosity_info()
logging.set_verbosity_debug()
LTXV_DEBUG=1
LTXV_FRAME_LOG_EVERY=8
import os, subprocess, shlex, tempfile
import torch
import json
import numpy as np
import random
import os
import shlex
import yaml
from typing import List, Dict
from pathlib import Path
import imageio
from PIL import Image
import tempfile
from huggingface_hub import hf_hub_download
import sys
import subprocess
import gc
import shutil
import contextlib
import time
import traceback
from einops import rearrange
import torch.nn.functional as F
from managers.vae_manager import vae_manager_singleton
from tools.video_encode_tool import video_encode_tool_singleton
DEPS_DIR = Path("/data")
LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"

def run_setup():
    setup_script_path = "setup.py"
    if not os.path.exists(setup_script_path):
        print("[DEBUG] 'setup.py' não encontrado. Pulando clonagem de dependências.")
        return
    try:
        print("[DEBUG] Executando setup.py para dependências...")
        subprocess.run([sys.executable, setup_script_path], check=True)
        print("[DEBUG] Setup concluído com sucesso.")
    except subprocess.CalledProcessError as e:
        print(f"[DEBUG] ERRO no setup.py (code {e.returncode}). Abortando.")
        sys.exit(1)
        
if not LTX_VIDEO_REPO_DIR.exists():
    print(f"[DEBUG] Repositório não encontrado em {LTX_VIDEO_REPO_DIR}. Rodando setup...")
    run_setup()

def add_deps_to_path():
    repo_path = str(LTX_VIDEO_REPO_DIR.resolve())
    if str(LTX_VIDEO_REPO_DIR.resolve()) not in sys.path:
        sys.path.insert(0, repo_path)
        print(f"[DEBUG] Repo adicionado ao sys.path: {repo_path}")

def _query_gpu_processes_via_nvml(device_index: int) -> List[Dict]:
    try:
        import psutil
        import pynvml as nvml
        nvml.nvmlInit()
        handle = nvml.nvmlDeviceGetHandleByIndex(device_index)
        try:
            procs = nvml.nvmlDeviceGetComputeRunningProcesses_v3(handle)
        except Exception:
            procs = nvml.nvmlDeviceGetComputeRunningProcesses(handle)
        results = []
        for p in procs:
            pid = int(p.pid)
            used_mb = None
            try:
                if getattr(p, "usedGpuMemory", None) is not None and p.usedGpuMemory not in (0,):
                    used_mb = max(0, int(p.usedGpuMemory) // (1024 * 1024))
            except Exception:
                used_mb = None
            name = "unknown"
            user = "unknown"
            try:
                import psutil
                pr = psutil.Process(pid)
                name = pr.name()
                user = pr.username()
            except Exception:
                pass
            results.append({"pid": pid, "name": name, "user": user, "used_mb": used_mb})
        nvml.nvmlShutdown()
        return results
    except Exception:
        return []

def _query_gpu_processes_via_nvidiasmi(device_index: int) -> List[Dict]:
    cmd = f"nvidia-smi -i {device_index} --query-compute-apps=pid,process_name,used_memory --format=csv,noheader,nounits"
    try:
        out = subprocess.check_output(shlex.split(cmd), stderr=subprocess.STDOUT, text=True, timeout=2.0)
    except Exception:
        return []
    results = []
    for line in out.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) >= 3:
            try:
                pid = int(parts[0]); name = parts[1]; used_mb = int(parts[2])
                user = "unknown"
                try:
                    import psutil
                    pr = psutil.Process(pid)
                    user = pr.username()
                except Exception:
                    pass
                results.append({"pid": pid, "name": name, "user": user, "used_mb": used_mb})
            except Exception:
                continue
    return results

def calculate_padding(orig_h, orig_w, target_h, target_w):
    pad_h = target_h - orig_h
    pad_w = target_w - orig_w
    pad_top = pad_h // 2
    pad_bottom = pad_h - pad_top
    pad_left = pad_w // 2
    pad_right = pad_w - pad_left
    return (pad_left, pad_right, pad_top, pad_bottom)

def calculate_new_dimensions(orig_w, orig_h, divisor=8):
    if orig_w == 0 or orig_h == 0:
        return 512, 512
    if orig_w >= orig_h:
        aspect_ratio = orig_w / orig_h
        new_h = 512
        new_w = new_h * aspect_ratio
    else:
        aspect_ratio = orig_h / orig_w
        new_w = 512
        new_h = new_w * aspect_ratio
    final_w = int(round(new_w / divisor)) * divisor
    final_h = int(round(new_h / divisor)) * divisor
    final_w = max(divisor, final_w)
    final_h = max(divisor, final_h)
    print(f"[Dimension Calc] Original: {orig_w}x{orig_h} -> Calculado: {new_w:.0f}x{new_h:.0f} -> Final (divisível por {divisor}): {final_w}x{final_h}")
    return final_h, final_w

def _gpu_process_table(processes: List[Dict], current_pid: int) -> str:
    if not processes:
        return "  - Processos ativos: (nenhum)\n"
    processes = sorted(processes, key=lambda x: (x.get("used_mb") or 0), reverse=True)
    lines = ["  - Processos ativos (PID | USER | NAME | VRAM MB):"]
    for p in processes:
        star = "*" if p["pid"] == current_pid else " "
        used_str = str(p["used_mb"]) if p.get("used_mb") is not None else "N/A"
        lines.append(f"    {star} {p['pid']} | {p['user']} | {p['name']} | {used_str}")
    return "\n".join(lines) + "\n"

def log_tensor_info(tensor, name="Tensor"):
    if not isinstance(tensor, torch.Tensor):
        print(f"\n[INFO] '{name}' não é tensor.")
        return
    print(f"\n--- Tensor: {name} ---")
    print(f"  - Shape: {tuple(tensor.shape)}")
    print(f"  - Dtype: {tensor.dtype}")
    print(f"  - Device: {tensor.device}")
    if tensor.numel() > 0:
        try:
            print(f"  - Min: {tensor.min().item():.4f}  Max: {tensor.max().item():.4f}  Mean: {tensor.mean().item():.4f}")
        except Exception:
            pass
    print("------------------------------------------\n")

add_deps_to_path()
from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem, LTXMultiScalePipeline
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
from ltx_video.models.autoencoders.vae_encode import un_normalize_latents, normalize_latents
from ltx_video.pipelines.pipeline_ltx_video import adain_filter_latent
from api.ltx.inference import (
    create_ltx_video_pipeline,
    create_latent_upsampler,
    load_image_to_tensor_with_resize_and_crop,
    seed_everething,
)

class VideoService:
    def __init__(self):
        t0 = time.perf_counter()
        print("[DEBUG] Inicializando VideoService...")
        self.debug = os.getenv("LTXV_DEBUG", "1") == "1"
        self.frame_log_every = int(os.getenv("LTXV_FRAME_LOG_EVERY", "8"))
        self.config = self._load_config()
        print(f"[DEBUG] Config carregada (precision={self.config.get('precision')}, sampler={self.config.get('sampler')})")
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"[DEBUG] Device selecionado: {self.device}")
        self.last_memory_reserved_mb = 0.0
        self._tmp_dirs = set(); self._tmp_files = set(); self._last_outputs = []

        self.pipeline, self.latent_upsampler = self._load_models()
        print(f"[DEBUG] Pipeline e Upsampler carregados. Upsampler ativo? {bool(self.latent_upsampler)}")

        print(f"[DEBUG] Movendo modelos para {self.device}...")
        self.pipeline.to(self.device)
        if self.latent_upsampler:
            self.latent_upsampler.to(self.device)

        self._apply_precision_policy()
        print(f"[DEBUG] runtime_autocast_dtype = {getattr(self, 'runtime_autocast_dtype', None)}")

        vae_manager_singleton.attach_pipeline(
            self.pipeline,
            device=self.device,
            autocast_dtype=self.runtime_autocast_dtype
        )
        print(f"[DEBUG] VAE manager conectado: has_vae={hasattr(self.pipeline, 'vae')} device={self.device}")

        if self.device == "cuda":
            torch.cuda.empty_cache()
            self._log_gpu_memory("Após carregar modelos")

        print(f"[DEBUG] VideoService pronto. boot_time={time.perf_counter()-t0:.3f}s")

    def _log_gpu_memory(self, stage_name: str):
        if self.device != "cuda":
            return
        device_index = torch.cuda.current_device() if torch.cuda.is_available() else 0
        current_reserved_b = torch.cuda.memory_reserved(device_index)
        current_reserved_mb = current_reserved_b / (1024 ** 2)
        total_memory_b = torch.cuda.get_device_properties(device_index).total_memory
        total_memory_mb = total_memory_b / (1024 ** 2)
        peak_reserved_mb = torch.cuda.max_memory_reserved(device_index) / (1024 ** 2)
        delta_mb = current_reserved_mb - getattr(self, "last_memory_reserved_mb", 0.0)
        processes = _query_gpu_processes_via_nvml(device_index) or _query_gpu_processes_via_nvidiasmi(device_index)
        print(f"\n--- [LOG GPU] {stage_name} (cuda:{device_index}) ---")
        print(f"  - Reservado: {current_reserved_mb:.2f} MB / {total_memory_mb:.2f} MB  (Δ={delta_mb:+.2f} MB)")
        if peak_reserved_mb > getattr(self, "last_memory_reserved_mb", 0.0):
            print(f"  - Pico reservado (nesta fase): {peak_reserved_mb:.2f} MB")
        print(_gpu_process_table(processes, os.getpid()), end="")
        print("--------------------------------------------------\n")
        self.last_memory_reserved_mb = current_reserved_mb

    def _register_tmp_dir(self, d: str):
        if d and os.path.isdir(d):
            self._tmp_dirs.add(d); print(f"[DEBUG] Registrado tmp dir: {d}")

    def _register_tmp_file(self, f: str):
        if f and os.path.exists(f):
            self._tmp_files.add(f); print(f"[DEBUG] Registrado tmp file: {f}")

    def finalize(self, keep_paths=None, extra_paths=None, clear_gpu=True):
        print("[DEBUG] Finalize: iniciando limpeza...")
        keep = set(keep_paths or []); extras = set(extra_paths or [])
        removed_files = 0
        for f in list(self._tmp_files | extras):
            try:
                if f not in keep and os.path.isfile(f):
                    os.remove(f); removed_files += 1; print(f"[DEBUG] Removido arquivo tmp: {f}")
            except Exception as e:
                print(f"[DEBUG] Falha removendo arquivo {f}: {e}")
            finally:
                self._tmp_files.discard(f)
        removed_dirs = 0
        for d in list(self._tmp_dirs):
            try:
                if d not in keep and os.path.isdir(d):
                    shutil.rmtree(d, ignore_errors=True); removed_dirs += 1; print(f"[DEBUG] Removido diretório tmp: {d}")
            except Exception as e:
                print(f"[DEBUG] Falha removendo diretório {d}: {e}")
            finally:
                self._tmp_dirs.discard(d)
        print(f"[DEBUG] Finalize: arquivos removidos={removed_files}, dirs removidos={removed_dirs}")
        gc.collect()
        try:
            if clear_gpu and torch.cuda.is_available():
                torch.cuda.empty_cache()
                try:
                    torch.cuda.ipc_collect()
                except Exception:
                    pass
        except Exception as e:
            print(f"[DEBUG] Finalize: limpeza GPU falhou: {e}")
        try:
            self._log_gpu_memory("Após finalize")
        except Exception as e:
            print(f"[DEBUG] Log GPU pós-finalize falhou: {e}")

    def _load_config(self):
        base = LTX_VIDEO_REPO_DIR / "configs"
        config_path = base / "ltxv-13b-0.9.8-distilled-fp8.yaml"
        print(f"[DEBUG] Carregando config: {config_path}")
        with open(config_path, "r") as file:
            return yaml.safe_load(file)

    def _load_models(self):
        t0 = time.perf_counter()
        LTX_REPO = "Lightricks/LTX-Video"
        print("[DEBUG] Baixando checkpoint principal...")
        distilled_model_path = hf_hub_download(repo_id=LTX_REPO, filename=self.config["checkpoint_path"])
        self.config["checkpoint_path"] = distilled_model_path
        print(f"[DEBUG] Checkpoint em: {distilled_model_path}")

        print("[DEBUG] Baixando upscaler espacial...")
        spatial_upscaler_path = hf_hub_download(repo_id=LTX_REPO, filename=self.config["spatial_upscaler_model_path"])
        self.config["spatial_upscaler_model_path"] = spatial_upscaler_path
        print(f"[DEBUG] Upscaler em: {spatial_upscaler_path}")

        print("[DEBUG] Construindo pipeline...")
        pipeline = create_ltx_video_pipeline(
            ckpt_path=self.config["checkpoint_path"],
            precision=self.config["precision"],
            text_encoder_model_name_or_path=self.config["text_encoder_model_name_or_path"],
            sampler=self.config["sampler"], device="cpu", enhance_prompt=False,
            prompt_enhancer_image_caption_model_name_or_path=self.config["prompt_enhancer_image_caption_model_name_or_path"],
            prompt_enhancer_llm_model_name_or_path=self.config["prompt_enhancer_llm_model_name_or_path"],
        )
        print("[DEBUG] Pipeline pronto.")

        latent_upsampler = None
        if self.config.get("spatial_upscaler_model_path"):
            print("[DEBUG] Construindo latent_upsampler...")
            latent_upsampler = create_latent_upsampler(self.config["spatial_upscaler_model_path"], device="cpu")
            print("[DEBUG] Upsampler pronto.")
        print(f"[DEBUG] _load_models() tempo total={time.perf_counter()-t0:.3f}s")
        return pipeline, latent_upsampler

    @torch.no_grad()
    def _upsample_latents_internal(self, latents: torch.Tensor) -> torch.Tensor:
        if not self.latent_upsampler:
            raise ValueError("Latent Upsampler não está carregado.")
        self.latent_upsampler.to(self.device)
        self.pipeline.vae.to(self.device)
        print(f"[DEBUG-UPSAMPLE] Shape de entrada: {tuple(latents.shape)}")
        latents_unnormalized = un_normalize_latents(latents, self.pipeline.vae, vae_per_channel_normalize=True)
        upsampled_latents = self.latent_upsampler(latents_unnormalized)
        upsampled_latents_normalized = normalize_latents(upsampled_latents, self.pipeline.vae, vae_per_channel_normalize=True)
        print(f"[DEBUG-UPSAMPLE] Shape de saída: {tuple(upsampled_latents_normalized.shape)}")
        return upsampled_latents_normalized
        
    def _apply_precision_policy(self):
        prec = str(self.config.get("precision", "")).lower()
        self.runtime_autocast_dtype = torch.float32
        print(f"[DEBUG] Aplicando política de precisão: {prec}")
        if prec in ["float8_e4m3fn", "bfloat16"]:
            self.runtime_autocast_dtype = torch.bfloat16
        elif prec == "mixed_precision":
            self.runtime_autocast_dtype = torch.float16

    def _prepare_conditioning_tensor(self, filepath, height, width, padding_values):
        print(f"[DEBUG] Carregando condicionamento: {filepath}")
        tensor = load_image_to_tensor_with_resize_and_crop(filepath, height, width)
        tensor = torch.nn.functional.pad(tensor, padding_values)
        out = tensor.to(self.device, dtype=self.runtime_autocast_dtype)
        print(f"[DEBUG] Cond shape={tuple(out.shape)} dtype={out.dtype} device={out.device}")
        return out

    def _concat_mp4s_no_reencode(self, mp4_list: List[str], out_path: str):
        if not mp4_list:
            raise ValueError("A lista de MP4s para concatenar está vazia.")
        if len(mp4_list) == 1:
            shutil.move(mp4_list[0], out_path)
            print(f"[DEBUG] Apenas um vídeo, movido para: {out_path}")
            return
            
        with tempfile.NamedTemporaryFile("w", delete=False, suffix=".txt") as f:
            for mp4 in mp4_list:
                f.write(f"file '{os.path.abspath(mp4)}'\n")
            list_path = f.name
    
        cmd = f"ffmpeg -y -f concat -safe 0 -i {list_path} -c copy {out_path}"
        print(f"[DEBUG] Concat: {cmd}")
    
        try:
            subprocess.check_call(shlex.split(cmd))
        finally:
            os.remove(list_path)

    def _save_and_log_video(self, pixel_tensor, base_filename, fps, temp_dir, results_dir, used_seed, progress_callback=None):
        """Função auxiliar para salvar um tensor de pixels em um arquivo MP4."""
        output_path = os.path.join(temp_dir, f"{base_filename}_{used_seed}.mp4")
        
        video_encode_tool_singleton.save_video_from_tensor(
            pixel_tensor, output_path, fps=fps, progress_callback=progress_callback
        )
        
        final_path = os.path.join(results_dir, f"{base_filename}_{used_seed}.mp4")
        shutil.move(output_path, final_path)
        print(f"[DEBUG] Vídeo salvo em: {final_path}")
        return final_path

    # ==============================================================================
    # --- NOVAS FUNÇÕES MODULARES ---
    # ==============================================================================

    def prepare_condition_items(self, items_list: List, height: int, width: int, num_frames: int):
        """
        Prepara a lista de tensores de condicionamento a partir de uma lista de imagens ou tensores.
        Formato da lista de entrada: [[media_path_ou_tensor, frame_alvo, peso], ...]
        """
        if not items_list:
            return []

        height_padded = ((height - 1) // 8 + 1) * 8
        width_padded = ((width - 1) // 8 + 1) * 8
        padding_values = calculate_padding(height, width, height_padded, width_padded)
        
        conditioning_items = []
        print("\n--- Preparando Itens de Condicionamento ---")
        for item in items_list:
            media, frame, weight = item
            
            if isinstance(media, str): 
                print(f"  - Carregando imagem: {media} para o frame {frame}")
                tensor = self._prepare_conditioning_tensor(media, height, width, padding_values)
            elif isinstance(media, torch.Tensor):
                print(f"  - Usando tensor fornecido para o frame {frame}")
                tensor = media.to(self.device, dtype=self.runtime_autocast_dtype)
            else:
                warnings.warn(f"Tipo de item desconhecido: {type(media)}. Ignorando.")
                continue

            safe_frame = max(0, min(int(frame), num_frames - 1))
            conditioning_items.append(ConditioningItem(tensor, safe_frame, float(weight)))

        print(f"Total de itens de condicionamento preparados: {len(conditioning_items)}")
        return conditioning_items

    def generate_low(self, prompt, negative_prompt, height, width, duration, guidance_scale, seed, conditioning_items=None):
        """
        Gera um vídeo em baixa resolução (primeiro passe).
        Retorna: (caminho_do_video_mp4, caminho_do_tensor_cpu, seed_usado)
        """
        print("\n--- INICIANDO ETAPA 1: GERAÇÃO EM BAIXA RESOLUÇÃO ---")
        self._log_gpu_memory("Início da Geração Low-Res")
        
        used_seed = random.randint(0, 2**32 - 1) if seed is None else int(seed)
        seed_everething(used_seed)
        
        FPS = 24.0
        target_frames = round(duration * FPS)
        actual_num_frames = max(9, int(round((target_frames - 1) / 8.0) * 8 + 1))
        
        height_padded = ((height - 1) // 8 + 1) * 8
        width_padded = ((width - 1) // 8 + 1) * 8
        generator = torch.Generator(device=self.device).manual_seed(used_seed)
        
        temp_dir = tempfile.mkdtemp(prefix="ltxv_low_"); self._register_tmp_dir(temp_dir)
        results_dir = "/app/output"; os.makedirs(results_dir, exist_ok=True)
        
        downscale_factor = self