import gradio as gr
from huggingface_hub import InferenceClient
import os
import logging
import asyncio
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
api_key = os.getenv("HUGGING_FACE_API_TOKEN")

# Available models
models = {
    "Llama-3B": "meta-llama/Llama-3.2-3B-Instruct",
    "DeepSeek-Coder": "deepseek-ai/deepseek-coder-1.3b-instruct",
    "DeepSeek-R1": "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
    "Mistral-7B": "mistralai/Mistral-7B-Instruct-v0.3"
}

# Initialize client
client = InferenceClient(api_key=api_key)

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Function to interact with selected model
async def chat_with_model(user_query, model_name):
    if model_name not in models:
        return "❌ Invalid model selection. Please choose a valid model."

    model_id = models[model_name]
    messages = [
        {"role": "system", "content": 
        """
        """
        },
        {"role": "user", "content": user_query}
    ]

    max_retries = 3
    for attempt in range(1, max_retries + 1):
        try:
            response = client.chat.completions.create(
                model=model_id,
                messages=messages,
                temperature=0.5,
                max_tokens=1024,  # Reduce for faster response
                top_p=0.7,
                stream=False
            )

            return response.choices[0].message.content
        except Exception as e:
            logger.warning(f"Attempt {attempt}/{max_retries} failed: {str(e)}")
            if attempt < max_retries:
                await asyncio.sleep(1)  # Short delay before retrying

    return "❌ The model is currently unavailable after multiple retries. Please try again later."

# Create Gradio UI
def chat_interface(user_query, model_name):
    return asyncio.run(chat_with_model(user_query, model_name))

with gr.Blocks() as demo:
    gr.Markdown("## Harry's AI Chatbot")
    gr.Markdown("### Select a model and ask your question to get a response from the AI.")
    
    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=list(models.keys()), 
            label="Select AI Model", 
            value="Mistral-7B"
        )
    
    user_input = gr.Textbox(label="Enter your message", placeholder="Type your question here...")
    chat_button = gr.Button("Chat")
    output_text = gr.Textbox(label="AI Response", interactive=False)

    chat_button.click(chat_interface, inputs=[user_input, model_dropdown], outputs=output_text)

# Launch Gradio app
demo.launch()