import os
import json
import uuid
import requests
import threading
import logging
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from google.cloud import storage
from google.auth import exceptions
from transformers import pipeline
from dotenv import load_dotenv
import uvicorn
import tempfile

load_dotenv()

API_KEY = os.getenv("API_KEY")
GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
GOOGLE_APPLICATION_CREDENTIALS_JSON = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON")
HF_API_TOKEN = os.getenv("HF_API_TOKEN")

logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

try:
    credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
    storage_client = storage.Client.from_service_account_info(credentials_info)
    bucket = storage_client.bucket(GCS_BUCKET_NAME)
    logger.info(f"Conexión con Google Cloud Storage exitosa. Bucket: {GCS_BUCKET_NAME}")
except (exceptions.DefaultCredentialsError, json.JSONDecodeError, KeyError, ValueError) as e:
    logger.error(f"Error al cargar las credenciales o bucket: {e}")
    raise RuntimeError(f"Error al cargar las credenciales o bucket: {e}")

app = FastAPI()

class DownloadModelRequest(BaseModel):
    model_name: str
    pipeline_task: str
    input_text: str

class GCSHandler:
    def __init__(self, bucket_name):
        self.bucket = storage_client.bucket(bucket_name)

    def file_exists(self, blob_name):
        exists = self.bucket.blob(blob_name).exists()
        logger.debug(f"Comprobando existencia de archivo '{blob_name}': {exists}")
        return exists

    def upload_file(self, blob_name, file_stream):
        blob = self.bucket.blob(blob_name)
        try:
            blob.upload_from_file(file_stream)
            logger.info(f"Archivo '{blob_name}' subido exitosamente a GCS.")
        except Exception as e:
            logger.error(f"Error subiendo el archivo '{blob_name}' a GCS: {e}")
            raise HTTPException(status_code=500, detail=f"Error subiendo archivo '{blob_name}' a GCS")

    def download_file(self, blob_name):
        blob = self.bucket.blob(blob_name)
        if not blob.exists():
            logger.error(f"Archivo '{blob_name}' no encontrado en GCS.")
            raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found.")
        logger.debug(f"Descargando archivo '{blob_name}' de GCS.")
        return blob.open("rb")

    def generate_signed_url(self, blob_name, expiration=3600):
        blob = self.bucket.blob(blob_name)
        url = blob.generate_signed_url(expiration=expiration)
        logger.debug(f"Generada URL firmada para '{blob_name}': {url}")
        return url

def download_model_from_huggingface(model_name):
    url = f"https://huggingface.co/{model_name}/tree/main"
    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
    try:
        logger.info(f"Descargando el modelo '{model_name}' desde Hugging Face...")
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            model_files = [
                "pytorch_model.bin",
                "config.json",
                "tokenizer.json",
                "model.safetensors",
            ]
            for file_name in model_files:
                file_url = f"https://huggingface.co/{model_name}/resolve/main/{file_name}"
                file_content = requests.get(file_url).content
                blob_name = f"{model_name}/{file_name}"
                blob = bucket.blob(blob_name)
                blob.upload_from_string(file_content)
                logger.info(f"Archivo '{file_name}' subido exitosamente al bucket GCS.")
        else:
            logger.error(f"Error al acceder al árbol de archivos de Hugging Face para '{model_name}'.")
            raise HTTPException(status_code=404, detail="Error al acceder al árbol de archivos de Hugging Face.")
    except Exception as e:
        logger.error(f"Error descargando archivos de Hugging Face: {e}")
        raise HTTPException(status_code=500, detail=f"Error descargando archivos de Hugging Face: {e}")

@app.post("/predict/")
async def predict(request: DownloadModelRequest):
    logger.info(f"Iniciando predicción para el modelo '{request.model_name}' con tarea '{request.pipeline_task}'...")
    try:
        gcs_handler = GCSHandler(GCS_BUCKET_NAME)
        model_prefix = request.model_name
        model_files = [
            "pytorch_model.bin",
            "config.json",
            "tokenizer.json",
            "model.safetensors",
        ]
        
        model_files_exist = all(gcs_handler.file_exists(f"{model_prefix}/{file}") for file in model_files)
        
        if not model_files_exist:
            logger.info(f"Modelos no encontrados en GCS, descargando '{model_prefix}' desde Hugging Face...")
            download_model_from_huggingface(model_prefix)
        
        model_files_streams = {}
        with tempfile.TemporaryDirectory() as temp_dir:
            for file in model_files:
                if gcs_handler.file_exists(f"{model_prefix}/{file}"):
                    file_path = os.path.join(temp_dir, file)
                    with open(file_path, "wb") as f:
                        gcs_handler.download_file(f"{model_prefix}/{file}").readinto(f)
                    model_files_streams[file] = file_path
            
            if not all(key in model_files_streams for key in ["config.json", "tokenizer.json", "pytorch_model.bin"]):
                logger.error(f"Faltan archivos necesarios para el modelo '{model_prefix}'.")
                raise HTTPException(status_code=500, detail="Required model files missing.")
            
            if request.pipeline_task in ["text-generation", "translation", "summarization"]:
                pipe = pipeline(request.pipeline_task, model=model_files_streams["pytorch_model.bin"], tokenizer=model_files_streams["tokenizer.json"])
                result = pipe(request.input_text)
                logger.info(f"Resultado generado para la tarea '{request.pipeline_task}': {result[0]}")
                return {"response": result[0]}

    except Exception as e:
        logger.error(f"Error en la predicción: {e}")
        raise HTTPException(status_code=500, detail=f"Error en la predicción: {e}")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)