import gradio as gr
import openai
from openai import OpenAI
import os

def setup_client():
    api_key = os.environ.get('OPENROUTER_API_KEY')
    if not api_key:
        raise ValueError("OPENROUTER_API_KEY environment variable not set")
    
    return OpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=api_key,
    )

def chat_with_grok(message, history):
    """Main chat function"""
    if not message:
        return history, ""
    
    try:
        client = setup_client()
        
        # Create the completion
        completion = client.chat.completions.create(
            model="x-ai/grok-4",
            messages=[
                {
                    "role": "user",
                    "content": message
                }
            ],
            max_tokens=1000,
            temperature=0.7
        )
        
        response = completion.choices[0].message.content
        
        # Add to history
        history.append((message, response))
        
        return history, ""
        
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        history.append((message, error_msg))
        return history, ""

def clear_chat():
    """Clear the chat history"""
    return [], ""

# Create the Gradio interface
with gr.Blocks(title="Grok 4 Chat Interface by Xhaheen ", theme=gr.themes.Soft()) as demo:
    gr.HTML("""
    <div style="text-align: center; padding: 20px;">
        <h1>🤖 Grok 4 Chat Interface</h1>
        <p>Chat with xAI's Grok 4 model via OpenRouter</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary", size="large")
                
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                label="Chat with Grok 4",
                height=500,
                show_copy_button=True
            )
            
            msg_input = gr.Textbox(
                label="Message",
                placeholder="Type your message here...",
                lines=2,
                max_lines=5
            )
            
            send_btn = gr.Button("Send 🚀", variant="primary", size="large")
    
    # Event handlers
    def submit_message(message, history):
        if message:
            return chat_with_grok(message, history)
        return history, message
    
    # Send message on button click
    send_btn.click(
        submit_message,
        inputs=[msg_input, chatbot],
        outputs=[chatbot, msg_input]
    )
    
    # Send message on Enter key
    msg_input.submit(
        submit_message,
        inputs=[msg_input, chatbot],
        outputs=[chatbot, msg_input]
    )
    
    # Clear chat
    clear_btn.click(
        clear_chat,
        outputs=[chatbot, msg_input]
    )

# Launch the interface - Remove share=True for Hugging Face
if __name__ == "__main__":
    demo.launch()