# api/gpu_manager.py

import os
import torch

class GPUManager:
    """
    Gerencia e aloca GPUs disponíveis para diferentes serviços (LTX, SeedVR).
    """
    def __init__(self):
        self.total_gpus = torch.cuda.device_count()
        self.ltx_gpus = []
        self.seedvr_gpus = []
        self._allocate_gpus()

    def _allocate_gpus(self):
        """
        Divide as GPUs disponíveis entre os serviços LTX e SeedVR.
        """
        print("="*50)
        print("🤖 Gerenciador de GPUs inicializado.")
        print(f"   > Total de GPUs detectadas: {self.total_gpus}")

        if self.total_gpus == 0:
            print("   > Nenhuma GPU detectada. Operando em modo CPU.")
        elif self.total_gpus == 1:
            print("   > 1 GPU detectada. Modo de compartilhamento de memória será usado.")
            # Ambos usarão a GPU 0, mas precisarão gerenciar a memória
            self.ltx_gpus = [0]
            self.seedvr_gpus = [0]
        else:
            # Divide as GPUs entre os dois serviços
            mid_point = self.total_gpus // 2
            self.ltx_gpus = list(range(0, mid_point))
            self.seedvr_gpus = list(range(mid_point, self.total_gpus))
            print(f"   > Alocação: LTX usará GPUs {self.ltx_gpus}, SeedVR usará GPUs {self.seedvr_gpus}.")
        
        print("="*50)

    def get_ltx_device(self):
        """Retorna o dispositivo principal para o LTX."""
        if not self.ltx_gpus:
            return torch.device("cpu")
        # Por padrão, o modelo principal do LTX roda na primeira GPU do seu grupo
        return torch.device(f"cuda:{self.ltx_gpus[0]}")

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

    def requires_memory_swap(self) -> bool:
        """Verifica se é necessário mover modelos entre CPU e GPU."""
        return self.total_gpus < 2

# Instância global para ser importada por outros módulos
gpu_manager = GPUManager()