import gradio as gr
import openai
import os
from typing import Optional

# System prompt for reformatting rough drafts into polished system prompts
SYSTEM_PROMPT = """You are a system prompt specialist that converts rough drafts into polished, clear specifications for AI code generation agents and subagents.

Your task:
1. Take the user's rough draft describing an AI code generation agent's desired functionalities
2. Transform it into a well-organized, clear system prompt written in second person ("you")
3. Ensure all desired functionalities from the original draft are preserved and clearly reflected
4. Focus on agents that will operate as subagents within multi-agent code generation frameworks

Requirements:
- Write the system prompt directing the agent in second person
- Organize the content logically with clear structure
- Maintain all functional requirements from the original draft
- Return only the polished system prompt in markdown within a code fence
- Treat each input as independent - do not reference previous specifications

Output format: Return the improved system prompt in markdown code fences with no additional commentary."""

def reformat_prompt(api_key: str, rough_draft: str) -> str:
    """
    Reformat a rough draft into a polished system prompt using OpenAI API.
    
    Args:
        api_key: OpenAI API key
        rough_draft: The rough draft text to be reformatted
        
    Returns:
        Reformatted system prompt or error message
    """
    if not api_key.strip():
        return "Error: Please provide your OpenAI API key."
    
    if not rough_draft.strip():
        return "Error: Please provide a rough draft to reformat."
    
    try:
        # Initialize OpenAI client with the provided API key
        client = openai.OpenAI(api_key=api_key)
        
        # Make API call to reformat the prompt
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": rough_draft}
            ],
            temperature=0.1,
            max_tokens=2000
        )
        
        reformatted_text = response.choices[0].message.content
        return reformatted_text
        
    except openai.AuthenticationError:
        return "Error: Invalid API key. Please check your OpenAI API key."
    except openai.RateLimitError:
        return "Error: Rate limit exceeded. Please try again later."
    except openai.APIError as e:
        return f"Error: OpenAI API error - {str(e)}"
    except Exception as e:
        return f"Error: {str(e)}"

def create_interface():
    """Create and configure the Gradio interface."""
    
    with gr.Blocks(
        title="Code Gen Prompt Reformatter",
        theme=gr.themes.Soft(),
        css="""
        .container {
            max-width: 1200px;
            margin: auto;
        }
        .header {
            text-align: center;
            margin-bottom: 2rem;
        }
        .api-key-box {
            background-color: #f8f9fa;
            padding: 1rem;
            border-radius: 8px;
            margin-bottom: 1rem;
        }
        .system-prompt-display {
            background-color: #f8f9fa;
            padding: 1rem;
            border-radius: 8px;
            font-family: monospace;
            white-space: pre-wrap;
            border: 1px solid #dee2e6;
        }
        """
    ) as interface:
        
        gr.HTML("""
        <div class="header">
            <h1>Code Gen Prompt Reformatter</h1>
            <p>Text transformation processor which attempts to improve a rough draft of a definition for a subagent in an agentic code generation system.</p>
        </div>
        """)
        
        with gr.Tabs():
            with gr.TabItem("Prompt Reformatter"):
                with gr.Row():
                    with gr.Column():
                        api_key_input = gr.Textbox(
                            label="OpenAI API Key",
                            placeholder="sk-...",
                            type="password",
                            info="Your API key is not stored and only used for this session"
                        )
                        
                        rough_draft_input = gr.Textbox(
                            label="Rough Draft",
                            placeholder="Enter your rough draft describing the desired functionalities of the AI code generation agent...",
                            lines=10,
                            max_lines=20
                        )
                        
                        with gr.Row():
                            clear_btn = gr.Button("Clear", variant="secondary")
                            submit_btn = gr.Button("Reformat Prompt", variant="primary", scale=2)
                    
                    with gr.Column():
                        output = gr.Textbox(
                            label="Reformatted System Prompt",
                            lines=15,
                            max_lines=25,
                            show_copy_button=True,
                            interactive=False
                        )
            
            with gr.TabItem("System Prompt"):
                gr.HTML("""
                <h3>System Prompt Used in Workflow</h3>
                <p>This is the system prompt that guides the AI in reformatting your rough drafts into polished system prompts:</p>
                """)
                
                gr.Textbox(
                    value=SYSTEM_PROMPT,
                    label="Current System Prompt",
                    lines=20,
                    max_lines=30,
                    show_copy_button=True,
                    interactive=False,
                    elem_classes=["system-prompt-display"]
                )
                
                gr.HTML("""
                <div style="margin-top: 1rem; padding: 1rem; background-color: #e3f2fd; border-radius: 8px;">
                    <h4>About This System Prompt</h4>
                    <p>This system prompt instructs the AI to:</p>
                    <ul>
                        <li>Convert rough notes into polished system prompts for code generation agents</li>
                        <li>Focus on sub-agents within multi-agent frameworks</li>
                        <li>Ensure all desired functionalities are clearly reflected</li>
                        <li>Return results in markdown format within code fences</li>
                        <li>Treat each input as a new, independent specification</li>
                    </ul>
                </div>
                """)
        
        # Event handlers
        submit_btn.click(
            fn=reformat_prompt,
            inputs=[api_key_input, rough_draft_input],
            outputs=output,
            show_progress=True
        )
        
        clear_btn.click(
            fn=lambda: ("", "", ""),
            outputs=[api_key_input, rough_draft_input, output]
        )
        
    
    return interface

# Create and launch the interface
if __name__ == "__main__":
    app = create_interface()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )