import gradio as gr
from openai import OpenAI

def generate_solution(api_key, problem_statement):
    try:
        client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=api_key.strip(),
        )

        completion = client.chat.completions.create(
            model="open-r1/olympiccoder-7b:free",
            messages=[
                {
                    "role": "system",
                    "content": "You are a competitive programming expert. Provide a correct solution with clear reasoning. First output the code, then explain the approach."
                },
                {
                    "role": "user",
                    "content": f"Solve this problem:\n{problem_statement}"
                }
            ],
            temperature=0.3,
            max_tokens=2048
        )

        response = completion.choices[0].message.content
        return format_response(response)
    
    except Exception as e:
        return f"Error: {str(e)}"

def format_response(response):
    # Split code and explanation
    formatted = []
    in_code = False
    code_buffer = []
    explanation_buffer = []
    
    for line in response.split('\n'):
        if line.strip().startswith('```'):
            in_code = not in_code
            if not in_code:
                formatted.append(f"```python\n{'\n'.join(code_buffer[1:])}\n```")
                code_buffer = []
            continue
            
        if in_code:
            code_buffer.append(line)
        else:
            explanation_buffer.append(line)
    
    if explanation_buffer:
        formatted.append("### Explanation\n" + '\n'.join(explanation_buffer))
        
    return "\n\n".join(formatted)

with gr.Blocks(title="Competitive Programming Assistant") as app:
    gr.Markdown("# 🏆 Competitive Programming Assistant")
    gr.Markdown("Powered by OlympicCoder-7B via OpenRouter AI")
    
    with gr.Row():
        api_key = gr.Textbox(
            label="OpenRouter API Key",
            type="password",
            placeholder="Enter your API key here..."
        )
    
    with gr.Row():
        problem_input = gr.Textbox(
            label="Problem Statement",
            lines=5,
            placeholder="Paste your programming problem here..."
        )
    
    submit_btn = gr.Button("Generate Solution", variant="primary")
    
    output = gr.Markdown()
    
    submit_btn.click(
        fn=generate_solution,
        inputs=[api_key, problem_input],
        outputs=output
    )

    gr.Examples(
        examples=[
            [
                "Given an array of integers, find two numbers such that they add up to a specific target number."
            ],
            [
                "Implement a function to calculate the minimum number of operations required to transform one string into another using only insertion, deletion, and substitution."
            ]
        ],
        inputs=[problem_input]
    )

if __name__ == "__main__":
    app.launch(server_port=7860, share=False)