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

def solve_competitive_problem(problem_statement, language_choice, api_key, progress=gr.Progress()):
    """
    Generate a solution for a competitive programming problem
    
    Args:
        problem_statement (str): The problem statement
        language_choice (str): Programming language for the solution
        api_key (str): OpenRouter API key
        progress: Gradio progress tracker
    
    Returns:
        str: Step-by-step solution with code
    """
    if not api_key.strip():
        return "Error: Please provide your OpenRouter API key."
    
    if not problem_statement.strip():
        return "Error: Please provide a problem statement."
    
    try:
        progress(0.1, "Initializing...")
        
        # Initialize OpenAI client with OpenRouter base URL
        client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=api_key,
        )
        
        progress(0.3, "Creating prompt...")
        
        # Create a more detailed prompt with language preference
        prompt = f"""
You are an expert competitive programmer. Analyze the following problem and provide a step-by-step solution with explanations and code in {language_choice}.
Problem:
{problem_statement}
Your response should include:
1. Problem Analysis:
   - Clear restatement of the problem in your own words
   - Identification of input/output formats
   - Key constraints and edge cases to consider
   - Time and space complexity requirements
2. Approach:
   - High-level strategy to solve the problem
   - Why this approach is optimal compared to alternatives
   - Any mathematical insights or observations
   - Data structures that will be helpful
3. Algorithm:
   - Detailed step-by-step breakdown of the algorithm
   - Clear explanation of the logic behind each step
   - Time complexity analysis with justification
   - Space complexity analysis with justification
   - Any optimizations made to improve performance
4. Implementation:
   - Clean, efficient, and well-commented {language_choice} code
   - Proper variable naming and code organization
   - Error handling and edge case management
   - Optimized for both readability and performance
5. Testing:
   - Example test cases with expected outputs
   - Edge case testing scenarios
   - Explanation of how to verify correctness
   - Potential areas where the solution might need improvement
Format your answer with clear headings and subheadings. Use markdown formatting for better readability.
        """
        
        progress(0.5, "Generating solution...")
        
        # Call the model
        completion = client.chat.completions.create(
            extra_headers={
                "HTTP-Referer": "https://competitive-programming-assistant.app",
                "X-Title": "Competitive Programming Assistant",
            },
            model="open-r1/olympiccoder-7b:free",
            messages=[
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            temperature=0.7,
            stream=False
        )
        
        progress(0.9, "Processing response...")
        
        solution = completion.choices[0].message.content
        
        progress(1.0, "Complete!")
        return solution
        
    except Exception as e:
        return f"Error: {str(e)}"

# Custom CSS for dark mode toggle and styling
custom_css = """
.dark-mode {
    --color-background: #1a1a1a;
    --color-text: #ffffff;
    --color-border: #444444;
    --color-accent: #3d85c6;
    --color-input-background: #2d2d2d;
    --color-button: #3d85c6;
    --color-button-hover: #5699d8;
}

.light-mode {
    --color-background: #ffffff;
    --color-text: #333333;
    --color-border: #dddddd;
    --color-accent: #2e67a9;
    --color-input-background: #f9f9f9;
    --color-button: #2e67a9;
    --color-button-hover: #4a7dbd;
}

body.dark-mode {
    background-color: var(--color-background);
    color: var(--color-text);
}

body.light-mode {
    background-color: var(--color-background);
    color: var(--color-text);
}

.theme-toggle {
    display: flex;
    align-items: center;
    justify-content: flex-end;
    padding: 8px 16px;
}

.theme-toggle button {
    background-color: var(--color-button);
    color: white;
    border: none;
    padding: 8px 16px;
    border-radius: 4px;
    cursor: pointer;
    transition: background-color 0.3s;
}

.theme-toggle button:hover {
    background-color: var(--color-button-hover);
}

.gradio-container {
    background-color: var(--color-background) !important;
    color: var(--color-text) !important;
}

.dark-mode .gradio-container textarea,
.dark-mode .gradio-container input {
    background-color: var(--color-input-background) !important;
    color: var(--color-text) !important;
    border-color: var(--color-border) !important;
}

.dark-mode .gradio-markdown {
    color: var(--color-text) !important;
}

.dark-mode .gradio-button {
    background-color: var(--color-button) !important;
}

.dark-mode .gradio-button:hover {
    background-color: var(--color-button-hover) !important;
}
"""

# JavaScript for theme toggle functionality
theme_toggle_js = """
function toggleTheme() {
    const body = document.body;
    if (body.classList.contains('light-mode')) {
        body.classList.remove('light-mode');
        body.classList.add('dark-mode');
        localStorage.setItem('theme', 'dark');
        document.getElementById('theme-toggle-btn').textContent = '☀️ Light Mode';
    } else {
        body.classList.remove('dark-mode');
        body.classList.add('light-mode');
        localStorage.setItem('theme', 'light');
        document.getElementById('theme-toggle-btn').textContent = '🌙 Dark Mode';
    }
}

document.addEventListener('DOMContentLoaded', function() {
    const savedTheme = localStorage.getItem('theme') || 'light';
    document.body.classList.add(savedTheme + '-mode');
    
    const themeToggleBtn = document.getElementById('theme-toggle-btn');
    if (themeToggleBtn) {
        themeToggleBtn.textContent = savedTheme === 'dark' ? '☀️ Light Mode' : '🌙 Dark Mode';
        themeToggleBtn.addEventListener('click', toggleTheme);
    }
});
"""

# Create a Gradio interface with theme support
with gr.Blocks(title="Competitive Programming Assistant", 
               theme=gr.themes.Soft(), 
               css=custom_css) as app:
    
    # Add JavaScript for theme toggle
    gr.HTML(f"""
    <script>
    {theme_toggle_js}
    </script>
    <div class="theme-toggle">
        <button id="theme-toggle-btn">🌙 Dark Mode</button>
    </div>
    """)
    
    gr.Markdown("""
    # 🏆 Competitive Programming Assistant
    
    Upload a problem statement from Codeforces, LeetCode, or any competitive programming platform to get:
    - Step-by-step analysis
    - Optimal solution approach
    - Complete code implementation
    - Time and space complexity analysis
    
    Powered by the OlympicCoder model.
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            api_key_input = gr.Textbox(
                placeholder="Enter your OpenRouter API key here",
                label="OpenRouter API Key",
                type="password"
            )
            
            problem_input = gr.Textbox(
                placeholder="Paste your competitive programming problem statement here...",
                label="Problem Statement",
                lines=10
            )
            
            language = gr.Dropdown(
                choices=["Python", "C++", "Java", "JavaScript"],
                value="Python",
                label="Programming Language"
            )
            
            submit_btn = gr.Button("Generate Solution", variant="primary")
        
        with gr.Column(scale=3):
            solution_output = gr.Markdown(
                label="Generated Solution"
            )
    
    with gr.Accordion("About", open=False):
        gr.Markdown("""
        ### How to use this app:
        
        1. Enter your OpenRouter API key (get one at [openrouter.ai](https://openrouter.ai))
        2. Paste the complete problem statement
        3. Select your preferred programming language
        4. Click "Generate Solution"
        
        ### Tips for best results:
        
        - Include the entire problem, including input/output formats and constraints
        - Make sure to include example inputs and outputs
        - For complex problems, consider adding clarifying notes
        
        ### Solution Format:
        
        Your solution will include:
        1. **Problem Analysis** - Breaking down the problem into manageable components
        2. **Approach** - Strategic methodology with mathematical insights
        3. **Algorithm** - Detailed step-by-step procedure with complexity analysis
        4. **Implementation** - Clean, commented code in your chosen language
        5. **Testing** - Example test cases and verification methods
        """)
    
    # Handle form submission
    submit_btn.click(
        solve_competitive_problem,
        inputs=[problem_input, language, api_key_input],
        outputs=solution_output
    )
    
    # Initialize theme based on user preference or system default
    gr.HTML("""
    <script>
        document.addEventListener('DOMContentLoaded', function() {
            const prefersDarkMode = window.matchMedia('(prefers-color-scheme: dark)').matches;
            const savedTheme = localStorage.getItem('theme');
            
            if (savedTheme) {
                document.body.classList.add(savedTheme + '-mode');
                if (savedTheme === 'dark') {
                    document.getElementById('theme-toggle-btn').textContent = '☀️ Light Mode';
                }
            } else if (prefersDarkMode) {
                document.body.classList.add('dark-mode');
                localStorage.setItem('theme', 'dark');
                document.getElementById('theme-toggle-btn').textContent = '☀️ Light Mode';
            } else {
                document.body.classList.add('light-mode');
                localStorage.setItem('theme', 'light');
            }
        });
    </script>
    """)

# Launch the app
if __name__ == "__main__":
    app.launch()