import gradio as gr
import random
import time
from typing import Dict, List, Tuple, Optional
import json

# Mock LLM UI generators (simulating different LLM responses)
def generate_ui_llm1(prompt: str, style: str = "modern") -> str:
    """Simulate LLM 1 generating UI code"""
    time.sleep(0.5)  # Simulate API call
    
    ui_templates = {
        "form": f"""
with gr.Blocks() as ui1:
    gr.Markdown("# {prompt.title()} - Generated by LLM 1")
    gr.Textbox(label="Name", placeholder="Enter your name")
    gr.Number(label="Age", minimum=1, maximum=120)
    gr.Dropdown(["Option A", "Option B", "Option C"], label="Choice")
    gr.Button("Submit", variant="primary")
""",
        "dashboard": f"""
with gr.Blocks() as ui1:
    gr.Markdown("## {prompt.title()} Dashboard - LLM 1")
    with gr.Row():
        with gr.Column():
            gr.Number(label="Metric 1", value=42)
            gr.Number(label="Metric 2", value=128)
        with gr.Column():
            gr.BarPlot(value={{"categories": ["A", "B", "C"], "values": [10, 20, 30]}})
    gr.Button("Refresh Data")
""",
        "chat": f"""
with gr.Blocks() as ui1:
    gr.Markdown("# {prompt.title()} Chat - LLM 1")
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Message", placeholder="Type your message...")
    gr.Button("Send")
""",
        "data": f"""
with gr.Blocks() as ui1:
    gr.Markdown("## {prompt.title()} Data View - LLM 1")
    gr.Dataframe(
        headers=["Name", "Score", "Status"],
        value=[["Alice", 95, "Active"], ["Bob", 87, "Pending"], ["Charlie", 92, "Active"]]
    )
    gr.Button("Export")
"""
    }
    
    # Choose template based on prompt keywords
    keywords = prompt.lower()
    if "form" in keywords or "input" in keywords:
        return ui_templates["form"]
    elif "dashboard" in keywords or "metric" in keywords or "chart" in keywords:
        return ui_templates["dashboard"]
    elif "chat" in keywords or "message" in keywords:
        return ui_templates["chat"]
    else:
        return ui_templates["data"]

def generate_ui_llm2(prompt: str, style: str = "minimal") -> str:
    """Simulate LLM 2 generating UI code"""
    time.sleep(0.5)  # Simulate API call
    
    ui_templates = {
        "form": f'''
with gr.Blocks() as ui2:
    gr.Markdown(f"### {prompt.title()} - LLM 2 Design")
    with gr.Group():
        gr.Textbox(label="Full Name", info="Enter your full name")
        gr.Slider(1, 100, value=25, label="Age Range")
        gr.Radio(["Choice 1", "Choice 2", "Choice 3"], label="Selection")
    with gr.Row():
        gr.Button("Submit", variant="primary")
        gr.Button("Reset")
''',
        "dashboard": f'''
with gr.Blocks() as ui2:
    gr.Markdown(f"## {prompt.title()} Analytics - LLM 2")
    with gr.Row():
        gr.LinePlot(value={{"x": [1,2,3,4], "y": [10,20,15,25]}})
        gr.ScatterPlot(value={{"x": [1,2,3,4], "y": [5,15,10,20]}})
    gr.Number(label="Total", value=100)
''',
        "chat": f'''
with gr.Blocks() as ui2:
    gr.Markdown(f"# {prompt.title()} Assistant - LLM 2")
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(height=400)
        with gr.Column(scale=1):
            gr.Dropdown(["Greeting", "Question", "Feedback"], label="Quick Actions")
    msg = gr.Textbox(label="Your input", placeholder="Ask anything...")
''',
        "data": f'''
with gr.Blocks() as ui2:
    gr.Markdown(f"## {prompt.title()} Explorer - LLM 2")
    gr.JSON(value={{"data": [{"id": 1, "name": "Item 1"}, {"id": 2, "name": "Item 2"}]}})
    gr.FileExplorer()
'''
    }
    
    # Choose template based on prompt keywords
    keywords = prompt.lower()
    if "form" in keywords or "input" in keywords:
        return ui_templates["form"]
    elif "dashboard" in keywords or "metric" in keywords or "chart" in keywords:
        return ui_templates["dashboard"]
    elif "chat" in keywords or "message" in keywords:
        return ui_templates["chat"]
    else:
        return ui_templates["data"]

class UIArena:
    def __init__(self):
        self.votes: Dict[str, int] = {"LLM 1": 0, "LLM 2": 0, "Tie": 0}
        self.current_comparison: Optional[Tuple[str, str]] = None
        
    def generate_comparison(self, prompt: str) -> Tuple[str, str, str]:
        """Generate UI code from both LLMs and update comparison"""
        ui1_code = generate_ui_llm1(prompt)
        ui2_code = generate_ui_llm2(prompt)
        
        self.current_comparison = (ui1_code, ui2_code)
        
        # Create visual previews (simplified for demonstration)
        preview1 = self.create_preview(ui1_code, "LLM 1")
        preview2 = self.create_preview(ui2_code, "LLM 2")
        
        return preview1, preview2, f"Comparison generated for: '{prompt}'"
    
    def create_preview(self, ui_code: str, llm_name: str) -> str:
        """Create a visual preview of the UI code"""
        return f"""
### {llm_name} Generated UI

{ui_code}
**Style Analysis:**
- Components: {ui_code.count('gr.') - ui_code.count('gr.Blocks')}
- Layout: {'Row & Column' if 'Row' in ui_code and 'Column' in ui_code else 'Single Column'}
- Variants: {'Yes' if 'variant=' in ui_code else 'No'}
"""
    
    def vote(self, choice: str) -> Tuple[str, str, str]:
        """Register a vote"""
        if choice in self.votes:
            self.votes[choice] += 1
        
        total_votes = sum(self.votes.values())
        if total_votes == 0:
            stats_text = "No votes yet"
        else:
            stats_text = "## Voting Results\n\n"
            for model, votes in self.votes.items():
                percentage = (votes / total_votes) * 100
                stats_text += f"**{model}**: {votes} votes ({percentage:.1f}%)\n"
        
        return f"Voted for {choice}!", stats_text, ""
    
    def get_stats(self) -> str:
        """Get current voting statistics"""
        total_votes = sum(self.votes.values())
        if total_votes == 0:
            return "No votes recorded yet. Generate a comparison and vote!"
        
        stats = "## 📊 UI Arena Statistics\n\n"
        for model, votes in self.votes.items():
            percentage = (votes / total_votes) * 100
            bar_length = int(percentage / 5)
            bar = "█" * bar_length + "░" * (20 - bar_length)
            stats += f"**{model}**: {votes} votes ({percentage:.1f}%)\n"
            stats += f"`{bar}`\n\n"
        
        stats += f"**Total Votes**: {total_votes}\n"
        
        if total_votes > 0:
            winner = max(self.votes, key=self.votes.get)
            if self.votes[winner] == max(self.votes.values()):
                stats += f"**Current Leader**: {winner}"
        
        return stats
    
    def reset_arena(self) -> Tuple[str, str, str, str]:
        """Reset the arena"""
        self.votes = {"LLM 1": 0, "LLM 2": 0, "Tie": 0}
        self.current_comparison = None
        return "", "", "Arena reset. Ready for new comparisons!", self.get_stats()

# Initialize arena
arena = UIArena()

def create_demo():
    """Create the main demo interface"""
    
    # Custom CSS for better layout
    custom_css = """
    .arena-container {
        border: 2px solid #e0e0e0;
        border-radius: 10px;
        padding: 20px;
        margin: 10px;
    }
    .comparison-box {
        min-height: 400px;
        border: 1px solid #d0d0d0;
        border-radius: 8px;
        padding: 15px;
        background: #f9f9f9;
    }
    .vote-button {
        margin: 5px;
        min-width: 120px;
    }
    """
    
    with gr.Blocks() as demo:
        gr.HTML("""
        <div style='text-align: center; margin-bottom: 20px;'>
            <h1>🏟️ UI Arena</h1>
            <p>Compare UIs generated by two different LLMs side by side</p>
            <p><a href='https://huggingface.co/spaces/akhaliq/anycoder' target='_blank'>Built with anycoder</a></p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                gr.Markdown("## 🎯 Generate Comparison")
                
                with gr.Row():
                    prompt_input = gr.Textbox(
                        label="UI Description Prompt",
                        placeholder="Describe the UI you want to generate (e.g., 'Create a user registration form', 'Build a metrics dashboard')",
                        lines=2
                    )
                
                with gr.Row():
                    generate_btn = gr.Button("🚀 Generate UIs", variant="primary", scale=2)
                    reset_btn = gr.Button("🔄 Reset Arena", variant="secondary", scale=1)
                
                status_output = gr.Textbox(label="Status", interactive=False)
            
            with gr.Column(scale=1):
                gr.Markdown("## 📈 Live Statistics")
                stats_display = gr.Markdown(arena.get_stats())
        
        gr.Markdown("---")
        
        with gr.Row():
            gr.Markdown("## 🔬 Side-by-Side Comparison", elem_classes="arena-container")
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("### LLM 1 Output", elem_id="llm1-header")
                ui1_output = gr.Markdown(
                    "Waiting for generation...",
                    elem_classes=["comparison-box"]
                )
                
            with gr.Column():
                gr.Markdown("### LLM 2 Output", elem_id="llm2-header")
                ui2_output = gr.Markdown(
                    "Waiting for generation...",
                    elem_classes=["comparison-box"]
                )
        
        with gr.Row():
            gr.Markdown("## 🗳️ Vote for Your Favorite")
        
        with gr.Row():
            vote_status = gr.Textbox(label="Vote Status", interactive=False)
        
        with gr.Row():
            with gr.Column():
                vote_llm1_btn = gr.Button("👍 Vote for LLM 1", variant="primary", elem_classes=["vote-button"])
            
            with gr.Column():
                vote_tie_btn = gr.Button("🤝 It's a Tie", variant="secondary", elem_classes=["vote-button"])
            
            with gr.Column():
                vote_llm2_btn = gr.Button("👍 Vote for LLM 2", variant="primary", elem_classes=["vote-button"])
        
        # Event handlers
        generate_btn.click(
            fn=arena.generate_comparison,
            inputs=[prompt_input],
            outputs=[ui1_output, ui2_output, status_output]
        ).then(
            fn=arena.get_stats,
            outputs=[stats_display]
        )
        
        reset_btn.click(
            fn=arena.reset_arena,
            outputs=[ui1_output, ui2_output, status_output, stats_display]
        )
        
        vote_llm1_btn.click(
            fn=lambda: arena.vote("LLM 1"),
            outputs=[vote_status, stats_display, ui1_output]
        )
        
        vote_llm2_btn.click(
            fn=lambda: arena.vote("LLM 2"),
            outputs=[vote_status, stats_display, ui2_output]
        )
        
        vote_tie_btn.click(
            fn=lambda: arena.vote("Tie"),
            outputs=[vote_status, stats_display, ui1_output]
        )
        
        # Auto-refresh stats periodically
        timer = gr.Timer(2)
        timer.tick(fn=arena.get_stats, outputs=[stats_display])
    
    return demo

# Create and launch the demo
if __name__ == "__main__":
    demo = create_demo()
    demo.launch(
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="indigo",
            neutral_hue="slate",
            text_size="lg",
            spacing_size="lg",
            radius_size="md"
        ),
        css="""
        .arena-container {
            border: 2px solid #e0e0e0;
            border-radius: 10px;
            padding: 20px;
            margin: 10px;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
        }
        .comparison-box {
            min-height: 400px;
            border: 1px solid #d0d0d0;
            border-radius: 8px;
            padding: 15px;
            background: #f8f9fa;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }
        .vote-button {
            margin: 5px;
            min-width: 120px;
            transition: all 0.3s ease;
        }
        .vote-button:hover {
            transform: translateY(-2px);
            box-shadow: 0 4px 8px rgba(0,0,0,0.2);
        }
        """,
        footer_links=[
            {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}
        ]
    )