from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import chromadb
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
import os
from openai import OpenAI
import zipfile

# Inicializar FastAPI
app = FastAPI()

# Configurar la API Key de OpenAI
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# Inicializar el cliente de OpenAI
client = OpenAI(api_key=OPENAI_API_KEY)

# Inicializar el cliente de ChromaDB en Hugging Face Space
chroma_client = chromadb.PersistentClient(path="chroma_db")

# Cargar la base de datos de Chroma como un vector store
vectorstore = Chroma(
    client=chroma_client,
    collection_name="docs",
    embedding_function=OpenAIEmbeddings(model="text-embedding-3-small", openai_api_key=OPENAI_API_KEY)
)

# Crear un retriever
retriever = vectorstore.as_retriever()

def obtener_extractos(pregunta: str):
    """Obtiene documentos relevantes desde ChromaDB"""
    docs_relevantes = retriever.invoke(pregunta)
    return [(doc.page_content, doc.metadata.get("url", "URL no disponible")) for doc in docs_relevantes]

# Modelo de datos para la solicitud
class ChatRequest(BaseModel):
    message: str
    system_message: str = "Eres un asistente virtual."
    max_tokens: int = 512
    temperature: float = 0.7
    top_p: float = 0.95

@app.post("/chat")
async def chat(request: ChatRequest):
    """Endpoint para generar respuestas usando OpenAI y ChromaDB"""
    
    # Obtener documentos relevantes
    contexto = obtener_extractos(request.message)

    # Construir el mensaje del sistema con el contexto
    system_message_final = f"""{request.system_message}
    Información relevante extraída de los documentos:
    {contexto}
    """

    messages = [
        {"role": "system", "content": system_message_final},
        {"role": "user", "content": request.message}
    ]
    print(messages)

    try:
        # Llamar a la API de OpenAI
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=messages,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            top_p=request.top_p
        )

        completion = response.choices[0].message.content
        return {"response": completion, "context": contexto}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Punto de entrada para ejecutar con Uvicorn en Hugging Face
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)