import gradio as gr
from huggingface_hub import InferenceClient
import time

# Initialize the client
client = InferenceClient("HuggingFaceH4/starchat2-15b-v0.1")

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    model_name
):
    """
    Generate chat responses using the specified model.
    """
    # Update client if model changes
    global client
    client = InferenceClient(model_name)
    
    messages = [{"role": "system", "content": system_message}]
    
    # Build conversation history
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    
    messages.append({"role": "user", "content": message})
    response = ""
    
    try:
        for message in client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            token = message.choices[0].delta.content
            response += token
            yield response
    except Exception as e:
        yield f"Error: {str(e)}"

def create_chat_interface():
    """
    Create and configure the Gradio interface
    """
    # Default system message
    default_system = """You are a helpful AI assistant. You provide accurate, informative, and engaging responses while being direct and concise."""
    
    # Available models
    models = [
        "HuggingFaceH4/starchat2-15b-v0.1",
        "meta-llama/Llama-2-70b-chat-hf",
        "mistralai/Mixtral-8x7B-Instruct-v0.1"
    ]

    # Create the interface
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🤖 Advanced AI Chatbot")
        
        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    height=600,
                    show_label=False,
                    container=True,
                    scale=2
                )
                msg = gr.Textbox(
                    show_label=False,
                    placeholder="Type your message here...",
                    container=False
                )
                
            with gr.Column(scale=1):
                with gr.Accordion("Settings", open=False):
                    system_msg = gr.Textbox(
                        label="System Message",
                        value=default_system,
                        lines=3
                    )
                    model = gr.Dropdown(
                        choices=models,
                        value=models[0],
                        label="Model"
                    )
                    max_tokens = gr.Slider(
                        minimum=50,
                        maximum=4096,
                        value=1024,
                        step=1,
                        label="Max Tokens"
                    )
                    temperature = gr.Slider(
                        minimum=0.1,
                        maximum=2.0,
                        value=0.7,
                        step=0.1,
                        label="Temperature"
                    )
                    top_p = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.9,
                        step=0.1,
                        label="Top P"
                    )
                
                with gr.Row():
                    clear = gr.Button("Clear Chat")
                    stop = gr.Button("Stop")
        
        # Handle sending messages
        msg.submit(
            respond,
            [msg, chatbot, system_msg, max_tokens, temperature, top_p, model],
            [chatbot],
            api_name="chat"
        )
        
        # Clear chat history
        clear.click(lambda: None, None, chatbot, queue=False)
        
        # Example prompts
        gr.Examples(
            examples=[
                ["Tell me a short story about a robot learning to paint."],
                ["Explain quantum computing in simple terms."],
                ["Write a haiku about artificial intelligence."]
            ],
            inputs=msg
        )

    return demo

# Create and launch the interface
if __name__ == "__main__":
    demo = create_chat_interface()
    demo.queue()
    demo.launch(share=True)