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)}"

# Create a Gradio interface
with gr.Blocks(title="Competitive Programming Assistant", theme=gr.themes.Soft()) as app:
    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
    )

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