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

def solve_competitive_problem(problem_statement, language_choice, api_key):
    """
    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
    
    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:
        # Initialize OpenAI client with OpenRouter base URL
        client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=api_key,
        )
        
        # Create 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: Break down the problem statement and identify the key challenges
2. Approach: Describe your approach to solve the problem and why it's optimal
3. Algorithm: Explain the algorithm you'll use with time and space complexity analysis
4. Implementation: Write clean, efficient, and well-commented {language_choice} code
5. Testing: Explain how to test the solution with example inputs/outputs
        """
        
        # 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
        )
        
        solution = completion.choices[0].message.content
        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
        """)
    
    # 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()