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

def solve_competitive_problem(problem_statement, language_choice, api_key, remember_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
        remember_api_key (bool): Whether to remember the API key
        progress: Gradio progress tracker
    
    Returns:
        str: Step-by-step solution with code
    """
    if not api_key or 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 with streaming enabled
        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=True  # Enable streaming for partial results
        )
        
        # Process streaming response
        solution = ""
        for chunk in completion:
            if chunk.choices[0].delta.content is not None:
                solution += chunk.choices[0].delta.content
                # Update progress and yield partial results
                progress(0.5 + (len(solution) % 500) / 1000, "Generating solution...")
                yield solution
        
        progress(1.0, "Complete!")
        return solution
        
    except Exception as e:
        return f"Error: {str(e)}"

def get_saved_api_key():
    """
    Return a JavaScript snippet that retrieves the saved API key from localStorage
    """
    return """
    function() {
        const savedKey = localStorage.getItem('openrouter_api_key');
        if (savedKey) {
            document.querySelector('#api_key_input input').value = savedKey;
            document.querySelector('#remember_api_key input').checked = true;
        }
    }
    """

def save_api_key(api_key, remember):
    """
    Return a JavaScript snippet that saves or removes the API key from localStorage
    """
    return f"""
    function() {{
        if ({str(remember).toLowerCase()}) {{
            localStorage.setItem('openrouter_api_key', "{api_key}");
        }} else {{
            localStorage.removeItem('openrouter_api_key');
        }}
    }}
    """

# Define theme options
light_theme = gr.themes.Soft()
dark_theme = gr.themes.Soft(
    primary_hue="slate",
    secondary_hue="slate",
    neutral_hue="slate",
    text_size=gr.themes.sizes.text_md,
)

# Create a Gradio interface
with gr.Blocks(title="Competitive Programming Assistant", theme=light_theme, css="""
    /* Custom CSS for better responsiveness and styling */
    @media (max-width: 768px) {
        .container { padding: 10px !important; }
        .responsive-padding { padding: 8px !important; }
    }
    .code-block { 
        background-color: #f7f7f7; 
        border-radius: 4px;
        padding: 10px;
        font-family: monospace;
    }
    .dark-mode .code-block {
        background-color: #2d2d2d;
        color: #f1f1f1;
    }
    /* Custom styling for buttons */
    .action-button {
        margin-top: 10px !important;
    }
""") as app:
    # JavaScript for theme toggling
    app.load(js="""
    function setupThemeToggle() {
        const themeToggle = document.getElementById('theme-toggle-btn');
        const savedTheme = localStorage.getItem('theme');
        
        if (savedTheme === 'dark') {
            document.body.classList.add('dark-mode');
            themeToggle.textContent = '☀️ Light Mode';
        }
        
        themeToggle.addEventListener('click', function() {
            document.body.classList.toggle('dark-mode');
            const isDark = document.body.classList.contains('dark-mode');
            localStorage.setItem('theme', isDark ? 'dark' : 'light');
            themeToggle.textContent = isDark ? '☀️ Light Mode' : '🌙 Dark Mode';
            
            // Update code block styling
            document.querySelectorAll('pre code').forEach(block => {
                if (isDark) {
                    block.style.backgroundColor = '#2d2d2d';
                    block.style.color = '#f1f1f1';
                } else {
                    block.style.backgroundColor = '#f7f7f7';
                    block.style.color = '#333';
                }
            });
        });
    }
    
    // Initialize syntax highlighting
    function setupPrism() {
        if (typeof Prism !== 'undefined') {
            Prism.highlightAll();
        }
    }
    
    // Apply saved API key on load
    document.addEventListener('DOMContentLoaded', function() {
        setupThemeToggle();
        const savedKey = localStorage.getItem('openrouter_api_key');
        if (savedKey) {
            document.querySelector('#api_key_input input').value = savedKey;
            document.querySelector('#remember_api_key input').checked = true;
        }
    });
    """)

    # App header with theme toggle
    with gr.Row(elem_classes="responsive-padding"):
        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.
        """)
        theme_btn = gr.Button("🌙 Dark Mode", elem_id="theme-toggle-btn")
    
    with gr.Row():
        with gr.Column(scale=2, elem_classes="responsive-padding"):
            with gr.Group():
                api_key_input = gr.Textbox(
                    placeholder="Enter your OpenRouter API key here",
                    label="OpenRouter API Key",
                    type="password",
                    elem_id="api_key_input"
                )
                
                remember_api_key = gr.Checkbox(
                    label="Remember my API key",
                    elem_id="remember_api_key"
                )
            
            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", elem_classes="action-button")
        
        with gr.Column(scale=3, elem_classes="responsive-padding"):
            solution_output = gr.Markdown(
                label="Generated Solution",
                elem_id="solution-output"
            )
    
    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
        """)
    
    # Load saved API key on page load
    app.load(fn=None, inputs=None, outputs=None, js=get_saved_api_key())
    
    # Handle API key saving when checkbox changes
    remember_api_key.change(
        fn=None,
        inputs=[api_key_input, remember_api_key],
        outputs=None,
        js=lambda api_key, remember: save_api_key(api_key, remember)
    )
    
    # Handle form submission with streaming
    submit_btn.click(
        solve_competitive_problem,
        inputs=[problem_input, language, api_key_input, remember_api_key],
        outputs=solution_output,
        js="""
        function(problem, language, api_key, remember) {
            // Before submission, save API key if requested
            if (remember) {
                localStorage.setItem('openrouter_api_key', api_key);
            }
            
            // Clear existing solution
            document.getElementById('solution-output').innerHTML = 'Generating solution...';
            
            // Enable syntax highlighting when solution loads
            setTimeout(function() {
                if (typeof Prism !== 'undefined') {
                    Prism.highlightAll();
                }
            }, 1000);
        }
        """
    )
    
    # Add JS for syntax highlighting
    app.load(js="""
    // Load Prism.js for syntax highlighting
    const prismScript = document.createElement('script');
    prismScript.src = 'https://cdnjs.cloudflare.com/ajax/libs/prism/1.25.0/prism.min.js';
    document.head.appendChild(prismScript);
    
    const prismCss = document.createElement('link');
    prismCss.rel = 'stylesheet';
    prismCss.href = 'https://cdnjs.cloudflare.com/ajax/libs/prism/1.25.0/themes/prism.min.css';
    document.head.appendChild(prismCss);
    
    // Load additional language support
    const languages = ['python', 'cpp', 'java', 'javascript'];
    languages.forEach(lang => {
        const script = document.createElement('script');
        script.src = `https://cdnjs.cloudflare.com/ajax/libs/prism/1.25.0/components/prism-${lang}.min.js`;
        document.head.appendChild(script);
    });
    
    // Observe DOM changes to apply syntax highlighting to new content
    const observer = new MutationObserver(mutations => {
        mutations.forEach(mutation => {
            if (mutation.addedNodes.length) {
                setTimeout(() => {
                    if (typeof Prism !== 'undefined') {
                        Prism.highlightAll();
                    }
                }, 100);
            }
        });
    });
    
    observer.observe(document.getElementById('solution-output'), { 
        childList: true,
        subtree: true
    });
    """)

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