import os
import litellm
from fastapi import FastAPI, HTTPException, Header, Depends
from pydantic import BaseModel
import gradio as gr

# Configuration du modèle
model_name = "gemini/gemini-1.5-pro"
api_key = os.environ.get("GEMINI_API_KEY")

# Configuration de LiteLLM
litellm.api_key = api_key

app = FastAPI()

class PromptRequest(BaseModel):
    prompt: str
    model: str = "gemini/gemini-1.5-pro"  # Valeur par défaut

# Prompt de base
base_prompt = """
Tu es un développeur logiciel expérimenté. Tu dois toujours répondre en français d'une manière réfléchie et concise. 
"""

# Clé API pour l'authentification
USER_API_KEY = os.environ.get("USER_API_KEY")  # Récupère la valeur de la variable d'environnement

async def verify_api_key(x_api_key: str = Header(...)):
    if x_api_key != USER_API_KEY:
        raise HTTPException(status_code=401, detail="Clé API invalide")

# Historique de la conversation
conversation_history = []

def chatbot_function(user_input):
    conversation_history.append({"role": "user", "content": user_input})

    response = litellm.completion(
        model=model_name,  # Utilisation du modèle par défaut pour le chatbot
        messages=conversation_history
    )

    conversation_history.append({"role": "assistant", "content": response.choices[0].message.content})

    return conversation_history

# Interface Gradio
iface = gr.Interface(
    fn=chatbot_function,
    inputs=gr.Textbox(lines=2, placeholder="Votre message..."),
    outputs="chatbot",
    title="Chatbot",
)


@app.post("/", dependencies=[Depends(verify_api_key)])
async def predict(request: PromptRequest):
    try:
        # Ajout du prompt de base à la requête de l'utilisateur
        full_prompt = base_prompt + request.prompt 

        response = litellm.completion(
            model=request.model, 
            messages=[{"role": "user", "content": full_prompt}]
        )
        return {"response": response.choices[0].message.content, "model": request.model}
    except Exception as e:
        return {"error": str(e)}

# Lancement de l'interface Gradio après le lancement de FastAPI
if __name__ == "__main__":
    import uvicorn
    import threading

    # Lancer FastAPI dans un thread séparé
    api_thread = threading.Thread(target=uvicorn.run, args=(app,), kwargs={"host": "0.0.0.0", "port": 7860})
    api_thread.start()

    # Lancer Gradio sur le port 7860
    iface.launch(server_port=7861, share=True)