import gradio as gr
import requests
import os
MODELS = [
    "Qwen/Qwen2.5-1.5B-Instruct",
    "Qwen/Qwen2.5-0.5B-Instruct",
    "microsoft/phi-2",
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
]

def generate_flutter_code(prompt, api_token=None):
    if not api_token:
        api_token = os.getenv("HF_TOKEN")

    if not api_token:
        return "❌ Hugging Face API token is required."

    enhanced_prompt = f"""
You are an expert Flutter/Dart developer.

Task:
Generate clean, production-ready Flutter/Dart code.

Requirements:
- Proper widget structure
- Clear state management
- Comments for complex logic
- Error handling
- Best practices

User request:
{prompt}

Return ONLY Dart code.
"""

    headers = {
        "Authorization": f"Bearer {api_token}",
        "Content-Type": "application/json"
    }

    payload = {
        "inputs": enhanced_prompt,
        "parameters": {
            "max_new_tokens": 1024,
            "temperature": 0.2,
            "top_p": 0.95,
            "do_sample": True
        }
    }

    for model_id in MODELS:
        try:
            api_url = f"https://api-inference.huggingface.co/models/{model_id}"
            response = requests.post(api_url, headers=headers, json=payload, timeout=40)

            if response.status_code == 200:
                data = response.json()
                if isinstance(data, list) and "generated_text" in data[0]:
                    return data[0]["generated_text"].strip()

            elif response.status_code in (503, 504):
                continue  # model loading → try next

            elif response.status_code == 429:
                return "⚠️ Rate limit reached. Try again later."

        except requests.exceptions.RequestException:
            continue

    return "❌ All models failed. Try again later."

# Gradio interface
def create_interface():
    with gr.Blocks(title="Flutter Code Generator") as interface:
        gr.HTML("""
        <div class="title">🚀 Flutter Code Generator</div>
        <div class="subtitle">Transform your ideas into Flutter/Dart code using AI</div>
        """)

        with gr.Row():
            with gr.Column(scale=2):
                prompt_input = gr.Textbox(
                    label="Describe your Flutter component or feature",
                    placeholder="e.g., Create a login form with email and password fields, validation, and submit button...",
                    lines=6
                )

                api_token_input = gr.Textbox(
                    label="Hugging Face API Token (optional)",
                    placeholder="Enter your HF token or set HF_TOKEN env variable",
                    type="password"
                )

                generate_btn = gr.Button(
                    "Generate Flutter Code ✨",
                    variant="primary",
                    size="lg"
                )

            with gr.Column(scale=3):
                output_code = gr.Code(
                    label="Generated Flutter/Dart Code",
                    lines=20
                )

        with gr.Accordion("How to use", open=False):
            gr.Markdown("""
            ### Getting Started
            1. **Describe your component**: Write a clear description of the Flutter component you want to create
            2. **API Token**: Get your token from [Hugging Face](https://huggingface.co/settings/tokens) or set `HF_TOKEN` environment variable
            3. **Generate**: Click the button to generate production-ready Flutter/Dart code

            ### Tips for Better Results
            - Be specific about the UI components, layout, and functionality
            - Mention state management needs (e.g., "using Provider" or "using setState")
            - Specify styling preferences (colors, themes, responsive design)
            - Include error handling requirements
            - Mention any specific Flutter packages to use

            ### Example Prompts
            - "Create a shopping cart component with add/remove items, quantity counter, and total price calculation"
            - "Build a weather app screen with current conditions, 7-day forecast, and location search using Riverpod"
            - "Design a user profile page with avatar, edit form, and save functionality"
            """)

        # Event handling
        generate_btn.click(
            fn=generate_flutter_code,
            inputs=[prompt_input, api_token_input],
            outputs=output_code,
            api_name="generate_code"
        )

        # Examples
        gr.Examples(
            examples=[
                ["Create a beautiful login screen with email/password fields, remember me checkbox, and social login buttons"],
                ["Build a product card widget with image, title, price, rating stars, and add to cart button"],
                ["Design a settings page with toggle switches, dropdown menus, and theme selector"],
                ["Create a chat message bubble component with sender info, timestamp, and message status"],
                ["Build a dashboard with statistics cards, charts, and navigation drawer"]
            ],
            inputs=prompt_input,
            label="Example prompts to try:"
        )

    return interface

# Launch the app
if __name__ == "__main__":
    interface = create_interface()
    interface.queue(max_size=10)  # Enable queuing for multiple users
    interface.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", 7860)),
        show_error=True,
        share=False,  # Set to True for public sharing on HF Spaces
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px;
            margin: auto;
        }
        .title {
            text-align: center;
            color: #2563eb;
            font-size: 2.5em;
            font-weight: bold;
            margin-bottom: 1em;
        }
        .subtitle {
            text-align: center;
            color: #64748b;
            font-size: 1.2em;
            margin-bottom: 2em;
        }
        """
    )