import os
import re
import requests
import json
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from google.auth import exceptions
from google.cloud import storage
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from io import BytesIO
from dotenv import load_dotenv
import uvicorn

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")

def validate_bucket_name(bucket_name):
    if not isinstance(bucket_name, str):
        raise ValueError("Bucket name must be a string.")
    if len(bucket_name) < 3 or len(bucket_name) > 63:
        raise ValueError("Bucket name must be between 3 and 63 characters long.")
    if not re.match(r"^[a-z0-9][a-z0-9\-\.]*[a-z0-9]$", bucket_name):
        raise ValueError(
            f"Invalid bucket name '{bucket_name}'. Bucket names must:"
            " - Use only lowercase letters, numbers, hyphens (-), and periods (.)"
            " - Start and end with a letter or number."
        )
    if "--" in bucket_name or ".." in bucket_name or ".-" in bucket_name or "-." in bucket_name:
        raise ValueError(
            f"Invalid bucket name '{bucket_name}'. Bucket names cannot contain consecutive periods, hyphens, or use '.-' or '-.'"
        )
    return bucket_name

try:
    GCS_BUCKET_NAME = validate_bucket_name(GCS_BUCKET_NAME)
    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)
except (exceptions.DefaultCredentialsError, json.JSONDecodeError, KeyError, ValueError) as e:
    print(f"Error al cargar credenciales o bucket: {e}")
    exit(1)

app = FastAPI()

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

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

    def file_exists(self, blob_name):
        return self.bucket.blob(blob_name).exists()

    def stream_file_from_gcs(self, blob_name):
        blob = self.bucket.blob(blob_name)
        if not blob.exists():
            raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found in GCS.")
        return blob.download_as_bytes()

    def upload_file_to_gcs(self, blob_name, data_stream):
        blob = self.bucket.blob(blob_name)
        blob.upload_from_file(data_stream)

    def ensure_bucket_structure(self, model_prefix):
        required_files = ["config.json", "tokenizer.json"]
        for filename in required_files:
            blob_name = f"{model_prefix}/{filename}"
            if not self.file_exists(blob_name):
                self.bucket.blob(blob_name).upload_from_string("{}", content_type="application/json")

    def stream_model_files(self, model_prefix, model_patterns):
        model_files = {}
        for pattern in model_patterns:
            blobs = list(self.bucket.list_blobs(prefix=f"{model_prefix}/"))
            for blob in blobs:
                if re.match(pattern, blob.name.split('/')[-1]):
                    model_files[blob.name.split('/')[-1]] = BytesIO(blob.download_as_bytes())
        return model_files

def download_model_from_huggingface(model_name):
    file_patterns = [
        "pytorch_model.bin",
        "model.safetensors",
        "config.json",
        "tokenizer.json",
    ]
    for i in range(1, 100):
        file_patterns.append(f"pytorch_model-{i:05}-of-{100:05}")
        file_patterns.append(f"model-{i:05}")
    for filename in file_patterns:
        url = f"https://huggingface.co/{model_name}/resolve/main/{filename}"
        headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
        try:
            response = requests.get(url, headers=headers, stream=True)
            if response.status_code == 200:
                blob_name = f"{model_name}/{filename}"
                blob = bucket.blob(blob_name)
                blob.upload_from_file(BytesIO(response.content))
        except Exception as e:
            pass

@app.post("/predict/")
async def predict(request: DownloadModelRequest):
    try:
        gcs_handler = GCSStreamHandler(GCS_BUCKET_NAME)
        model_prefix = request.model_name
        model_patterns = [
            r"pytorch_model-\d+-of-\d+",
            r"model-\d+",
            r"pytorch_model.bin",
            r"model.safetensors",
        ]
        if not any(
            gcs_handler.file_exists(f"{model_prefix}/{pattern}") for pattern in model_patterns
        ):
            download_model_from_huggingface(model_prefix)
        model_files = gcs_handler.stream_model_files(model_prefix, model_patterns)
        config_stream = gcs_handler.stream_file_from_gcs(f"{model_prefix}/config.json")
        tokenizer_stream = gcs_handler.stream_file_from_gcs(f"{model_prefix}/tokenizer.json")
        model = AutoModelForCausalLM.from_pretrained(BytesIO(config_stream))
        tokenizer = AutoTokenizer.from_pretrained(BytesIO(tokenizer_stream))
        pipeline_task = request.pipeline_task
        pipeline_ = pipeline(pipeline_task, model=model, tokenizer=tokenizer)
        input_text = request.input_text
        result = pipeline_(input_text)
        return {"response": result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {e}")

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