import os
import gradio as gr
from huggingface_hub import InferenceClient

# Constants
DEFAULT_MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.1"
HF_TOKEN = os.environ.get("HF_TOKEN")

def chat_with_model(system_prompt, user_message, max_tokens=500, model_id=""):
    """
    Generate a chat completion using the specified or default Mistral model.
    
    Args:
    system_prompt (str): The system prompt to set the context.
    user_message (str): The user's input message.
    max_tokens (int): Maximum number of tokens to generate.
    model_id (str): The model ID to use. If empty, uses the default model.
    
    Returns:
    str: The model's response.
    """
    model_id = model_id.strip() if model_id else DEFAULT_MODEL_ID
    
    client = InferenceClient(model_id, token=HF_TOKEN)
    
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_message}
    ]
    
    response = ""
    try:
        for message in client.chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            stream=True
        ):
            response += message.choices[0].delta.content or ""
    except Exception as e:
        response = f"Error: {str(e)}"
    
    return response

def create_gradio_interface():
    """Create and configure the Gradio interface."""
    return gr.Interface(
        fn=chat_with_model,
        inputs=[
            gr.Textbox(label="System Prompt", placeholder="Enter the system prompt here..."),
            gr.Textbox(label="User Message", placeholder="Ask a question..."),
            gr.Slider(minimum=50, maximum=1000, value=500, step=50, label="Max Tokens"),
            gr.Textbox(label="Model ID", placeholder=f"Enter model ID (default: {DEFAULT_MODEL_ID})")
        ],
        outputs=gr.Textbox(label="Response"),
        title="Mistral Chatbot",
        description="Chat with Mistral model using your own system prompts and choose your model.",
        examples=[
            ["You are a helpful AI assistant.", "What is the capital of France?", 500, ""],
            ["You are an expert in Python programming.", "Explain list comprehensions.", 500, ""]
        ]
    )

if __name__ == "__main__":
    iface = create_gradio_interface()
    iface.launch(show_api=True, share=False, show_error=True)