# api/gpu_manager.py (Versão com VAE Dedicado)

import os
import torch
import math

class GPUManager:
    """
    Gerencia e aloca GPUs de forma inteligente entre LTX (com VAE dedicado),
    SeedVR e VINCIE.
    """
    def __init__(self):
        self.total_gpus = torch.cuda.device_count()
        self.ltx_main_gpus = []
        self.ltx_vae_gpu = []
        self.seedvr_gpus = []
        self.vincie_gpus = []
        self._allocate_gpus()

    def _allocate_gpus(self):
        print("="*60)
        print("🤖 Gerenciador de GPUs (com VAE dedicado para LTX)")
        print(f"   > Total de GPUs detectadas: {self.total_gpus}")

        all_indices = list(range(self.total_gpus))

        if self.total_gpus == 0:
            print("   > Nenhuma GPU detectada. Operando em modo CPU.")
        elif self.total_gpus == 1:
            print("   > 1 GPU: Modo de compartilhamento total.")
            self.ltx_main_gpus = [0]
            self.ltx_vae_gpu = [0] # Compartilha com o principal
            self.seedvr_gpus = [0]
            self.vincie_gpus = [0]
        elif self.total_gpus == 2:
            print("   > 2 GPUs: LTX com VAE dedicado, outros compartilham a GPU principal.")
            self.ltx_main_gpus = [0]
            self.ltx_vae_gpu = [1] # VAE fica com a segunda GPU
            self.seedvr_gpus = [0] # Compartilha com LTX principal
            self.vincie_gpus = [0] # Compartilha com LTX principal
        else: # 3 ou mais GPUs
            print(f"   > {self.total_gpus} GPUs: Alocação distribuída.")
            # LTX sempre fica com as duas primeiras GPUs se disponíveis
            self.ltx_main_gpus = [0]
            self.ltx_vae_gpu = [1]
            
            remaining_gpus = all_indices[2:]
            
            if not remaining_gpus: # Caso de exatamente 2 GPUs, já coberto, mas por segurança
                self.seedvr_gpus = [0]
                self.vincie_gpus = [0]
            else:
                # O resto é dividido entre SeedVR e VINCIE
                vincie_count = max(1, math.ceil(len(remaining_gpus) / 2))
                seedvr_count = len(remaining_gpus) - vincie_count
                
                self.vincie_gpus = remaining_gpus[:vincie_count]
                self.seedvr_gpus = remaining_gpus[vincie_count:]
        
        print(f"   > Alocação Final:")
        print(f"     - LTX (Transformer): GPUs {self.ltx_main_gpus}")
        print(f"     - LTX (VAE):         GPUs {self.ltx_vae_gpu}")
        print(f"     - SeedVR:            GPUs {self.seedvr_gpus}")
        print(f"     - VINCIE:            GPUs {self.vincie_gpus}")
        print("="*60)

    def get_ltx_device(self) -> torch.device:
        """Retorna o dispositivo principal para o Transformer do LTX."""
        if not self.ltx_main_gpus:
            return torch.device("cpu")
        return torch.device(f"cuda:{self.ltx_main_gpus[0]}")

    def get_ltx_vae_device(self) -> torch.device:
        """Retorna o dispositivo dedicado para o VAE do LTX."""
        if not self.ltx_vae_gpu:
            return torch.device("cpu")
        return torch.device(f"cuda:{self.ltx_vae_gpu[0]}")

    def get_seedvr_devices(self) -> list:
        """Retorna a lista de IDs de GPU para o SeedVR."""
        return self.seedvr_gpus
    
    def get_vincie_devices(self) -> list:
        """Retorna a lista de IDs de GPU para o VINCIE."""
        return self.vincie_gpus

    def requires_memory_swap(self) -> bool:
        """Verifica se múltiplos serviços estão compartilhando a mesma GPU."""
        all_allocations = self.ltx_main_gpus + self.seedvr_gpus + self.vincie_gpus
        # O swap é necessário se o número de alocações for maior que o número de GPUs únicas
        return len(all_allocations) > len(set(all_allocations))

# Instância global
gpu_manager = GPUManager()