import gradio as gr
import os
import requests
import time

# Get API key from environment
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"

# ✅ YOUR AVAILABLE MODELS (from your account)
AVAILABLE_MODELS = {
    "Llama 3.1 (8B) - Fast": "llama-3.1-8b-instant",
    "Llama 3.3 (70B) - Powerful": "llama-3.3-70b-versatile", 
    "Qwen 3 (32B)": "qwen/qwen3-32b",
    "Llama 4 Maverick (17B)": "meta-llama/llama-4-maverick-17b-128e-instruct",
    "Llama 4 Scout (17B)": "meta-llama/llama-4-scout-17b-16e-instruct",
    "GPT OSS (20B)": "openai/gpt-oss-20b",
    "GPT OSS (120B)": "openai/gpt-oss-120b"
}

# Default model that should definitely work
DEFAULT_MODEL = "llama-3.1-8b-instant"

def query_groq_api(message, chat_history, model_name):
    """Call GROQ API with your available models"""
    
    if not GROQ_API_KEY:
        return "❌ **ERROR**: GROQ_API_KEY not found. Add it in Hugging Face Secrets."
    
    # Get actual model ID from friendly name
    actual_model = AVAILABLE_MODELS.get(model_name, DEFAULT_MODEL)
    
    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }
    
    # Build messages array
    messages = [
        {
            "role": "system", 
            "content": "You are CodeMentor, a helpful programming tutor. Keep answers clear and concise (2-3 sentences). Always format code examples in markdown code blocks."
        }
    ]
    
    # Add conversation history (last 3 exchanges for context)
    for user_msg, bot_msg in chat_history[-3:]:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    
    # Add current message
    messages.append({"role": "user", "content": message})
    
    # Request data
    data = {
        "model": actual_model,
        "messages": messages,
        "temperature": 0.7,
        "max_tokens": 300,
        "top_p": 0.9
    }
    
    try:
        # Make API call
        response = requests.post(
            GROQ_API_URL,
            headers=headers,
            json=data,
            timeout=15
        )
        
        if response.status_code == 200:
            result = response.json()
            reply = result["choices"][0]["message"]["content"]
            return f"{reply}\n\n**🤖 Model**: {model_name}"
        elif response.status_code == 429:
            return "⏰ **RATE LIMITED**: Please wait 60 seconds before trying again."
        else:
            return f"❌ **API Error {response.status_code}**: Try a different model from the dropdown."
    
    except requests.exceptions.Timeout:
        return "⏱️ **TIMEOUT**: Request took too long. Try the 'Llama 3.1 (8B)' model."
    except Exception as e:
        return f"⚠️ **Error**: {str(e)[:100]}"

def chat_response(message, history, model):
    """Main chat function"""
    if not message.strip():
        return "", history
    
    # Get response from GROQ
    bot_message = query_groq_api(message, history, model)
    
    # Add to history
    history.append((message, bot_message))
    
    return "", history

def clear_chat():
    return []

def reset_settings():
    return ["Llama 3.1 (8B) - Fast", "What is Python?"]

# Create Gradio interface
with gr.Blocks(title="Programming Tutor - GROQ", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 👨‍💻 Programming Tutor Chatbot
    **Using YOUR Available GROQ Models**
    """)
    
    # Chatbot display
    chatbot = gr.Chatbot(height=400, label="Conversation")
    
    # Store chat history
    chat_state = gr.State([])
    
    with gr.Row():
        with gr.Column(scale=1):
            # Model selection dropdown
            gr.Markdown("### ⚙️ Settings")
            
            model_dropdown = gr.Dropdown(
                choices=list(AVAILABLE_MODELS.keys()),
                value="Llama 3.1 (8B) - Fast",
                label="Select Model",
                info="Choose from your available models"
            )
            
            # Quick actions
            gr.Markdown("### ⚡ Quick Actions")
            with gr.Row():
                clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary")
                reset_btn = gr.Button("🔄 Reset", variant="secondary")
            
            # Example questions
            gr.Markdown("### 💡 Example Questions")
            
            examples = [
                "Hello! Who are you?",
                "Explain variables in programming",
                "Write a Python function to add two numbers",
                "What are data types?"
            ]
            
            example_output = gr.Textbox(label="Selected Question", interactive=False, visible=False)
            
            for example in examples:
                btn = gr.Button(example, size="sm")
                btn.click(
                    lambda x=example: x,
                    inputs=None,
                    outputs=example_output
                )
        
        with gr.Column(scale=2):
            # Message input
            msg = gr.Textbox(
                placeholder="Type your programming question here...",
                label="Your Question",
                lines=3
            )
            
            # Send button
            send_btn = gr.Button("🚀 Send Message", variant="primary", size="lg")
    
    # Update message box when example is clicked
    example_output.change(
        lambda x: x,
        inputs=example_output,
        outputs=msg
    )
    
    # Function to handle message submission
    def respond(message, history, model):
        if not message.strip():
            return "", history
        
        bot_reply = query_groq_api(message, history, model)
        history.append((message, bot_reply))
        return "", history
    
    # Set up event handlers
    msg.submit(
        respond,
        inputs=[msg, chat_state, model_dropdown],
        outputs=[msg, chatbot]
    )
    
    send_btn.click(
        respond,
        inputs=[msg, chat_state, model_dropdown],
        outputs=[msg, chatbot]
    )
    
    clear_btn.click(
        lambda: ([], []),
        inputs=None,
        outputs=[chatbot, chat_state]
    )
    
    reset_btn.click(
        reset_settings,
        inputs=None,
        outputs=[model_dropdown, msg]
    )
    
    # Footer
    gr.Markdown(f"""
    ---
    ### ℹ️ About Your Available Models
    
    **Fast Models (Try These First):**
    1. **Llama 3.1 (8B)** - Fastest, good for simple questions
    2. **Qwen 3 (32B)** - Balanced speed and quality
    3. **Llama 4 Maverick (17B)** - Good for code examples
    
    **Powerful Models (Slower):**
    1. **Llama 3.3 (70B)** - Most powerful but slower
    2. **GPT OSS (120B)** - Very powerful, may be slow
    
    **Rate Limits:** 30 requests/minute for most models
    **Tip:** Start with "Llama 3.1 (8B)" for fastest responses
    """)

if __name__ == "__main__":
    demo.launch(debug=False, server_name="0.0.0.0", server_port=7860)