import gradio as gr
from openai import OpenAI

def generate_solution(api_key, problem_statement):
    # Check if inputs are provided
    if not api_key or not problem_statement:
        return "Please provide both an API key and a 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):
    formatted = []
    in_code = False
    code_buffer = []
    explanation_buffer = []
    
    for line in response.split('\n'):
        if line.strip().startswith('```'):
            if in_code and code_buffer:
                # Skip language marker if present
                if len(code_buffer) > 0 and any(code_buffer[0].strip().lower() == lang for lang in ["python", "java", "c++", "cpp"]):
                    code_content = '\n'.join(code_buffer[1:])
                else:
                    code_content = '\n'.join(code_buffer)
                formatted.append(f"```python\n{code_content}\n```")
                code_buffer = []
            in_code = not in_code
            continue
            
        if in_code:
            code_buffer.append(line)
        else:
            explanation_buffer.append(line)
    
    # Handle case where code block wasn't properly closed
    if in_code and code_buffer:
        if len(code_buffer) > 0 and any(code_buffer[0].strip().lower() == lang for lang in ["python", "java", "c++", "cpp"]):
            code_content = '\n'.join(code_buffer[1:])
        else:
            code_content = '\n'.join(code_buffer)
        formatted.append(f"```python\n{code_content}\n```")
    
    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...",
            value=""  # Set default value to empty string
        )
    
    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(label="Solution")
    
    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=True)  # Changed share to True to make it accessible externally