import json
from transformers import pipeline
import torch

class EndpointHandler:
    """
    Handler customizado para o endpoint de inferência.
    """
    def __init__(self, path=""):
        # O 'path' é o diretório raiz. Construímos o caminho completo para a pasta do modelo final.
        model_path = f"{path}/meu-modelo-ham10000-final"
        
        # Determina se usará GPU ou CPU
        device = 0 if torch.cuda.is_available() else -1
        
        self.pipeline = pipeline(
            "image-classification",
            model=model_path, # <--- AGORA APONTAMOS PARA A PASTA EXATA!
            device=device
        )
        print(f"Pipeline de classificação de imagem carregada com sucesso a partir de: {model_path}")

    def __call__(self, data: dict) -> list:
        """
        Esta função é chamada a cada requisição para a API.
        'data' é o payload da requisição.
        """
        # Extrai os inputs do payload da requisição
        inputs = data.pop("inputs", data)
        parameters = data.pop("parameters", None)

        print(f"Recebida requisição para analisar: {inputs}")

        # Roda a inferência
        try:
            if parameters:
                result = self.pipeline(inputs, **parameters)
            else:
                result = self.pipeline(inputs)

            print(f"Resultado da análise: {result}")
            return result
        except Exception as e:
            # Retorna um erro formatado se algo der errado durante a análise
            error_message = f"Erro ao processar a requisição: {str(e)}"
            print(error_message)
            return [{"error": error_message}]