import gradio as gr
from dotenv import load_dotenv
import os
from proctor import (
    CompositeTechnique,
    RolePrompting,
    ChainOfThought,
    ChainOfVerification,
    SelfAsk,
    EmotionPrompting,
    ZeroShotCoT,
    list_techniques,
)

# Load environment variables
load_dotenv()

# Check for OpenRouter API key
openrouter_key = os.environ.get("OPENROUTER_API_KEY")
if not openrouter_key:
    raise ValueError("OPENROUTER_API_KEY not set. Please set it in your .env file.")

# Available models and techniques
MODELS = {
    "Google Gemini 2.5 Flash": "openrouter/google/gemini-2.5-flash-preview-05-20",
    "Claude 4 Sonnet": "openrouter/anthropic/claude-sonnet-4",
    "DeepSeek R1": "openrouter/deepseek/deepseek-r1-0528",
    "Llama 4 Scout": "openrouter/meta-llama/llama-4-scout",
    "Mistral Small 3.1 24B": "openrouter/mistralai/mistral-small-3.1-24b-instruct",
}
TECHNIQUES = list_techniques()

# Model configurations
MODEL_CONFIGS = {
    "openrouter/google/gemini-2.5-flash-preview-05-20": {
        "api_base": "https://openrouter.ai/api/v1",
        "api_key": openrouter_key,
        "temperature": 0.3,
        "max_tokens": 1500,
    },
    "openrouter/anthropic/claude-sonnet-4": {
        "api_base": "https://openrouter.ai/api/v1",
        "api_key": openrouter_key,
        "temperature": 0.7,
        "max_tokens": 2000,
    },
    "openrouter/deepseek/deepseek-r1-0528": {
        "api_base": "https://openrouter.ai/api/v1",
        "api_key": openrouter_key,
        "temperature": 0.6,
        "max_tokens": 3000,
    },
    "openrouter/meta-llama/llama-4-scout": {
        "api_base": "https://openrouter.ai/api/v1",
        "api_key": openrouter_key,
        "temperature": 0.6,
        "max_tokens": 2500,
    },
    "openrouter/mistralai/mistral-small-3.1-24b-instruct": {
        "api_base": "https://openrouter.ai/api/v1",
        "api_key": openrouter_key,
        "temperature": 0.8,
        "max_tokens": 1000,
    },
}

# Composite technique definitions
TECHNIQUE_CONFIGS = {
    "Expert Chain-of-Thought": CompositeTechnique(
        name="Expert Chain-of-Thought",
        identifier="custom-expert-cot",
        techniques=[RolePrompting(), ChainOfThought(), ChainOfVerification()],
    ),
    "Deep Reasoning Analysis": CompositeTechnique(
        name="Deep Reasoning Analysis",
        identifier="deep-reasoning",
        techniques=[ChainOfThought(), SelfAsk(), ChainOfVerification()],
    ),
    "ChainOfThought": ChainOfThought(),
    "EmotionPrompting": EmotionPrompting(),
    "RolePrompting": RolePrompting(),
    "SelfAsk": SelfAsk(),
    "ZeroShotCoT": ZeroShotCoT(),
}

def format_as_markdown(response):
    """
    Format the response as Markdown for better readability.
    Assumes the response may have sections like headings, lists, etc.
    """
    lines = response.split("\n")
    formatted_lines = []
    in_list = False

    for line in lines:
        line = line.strip()
        if not line:
            in_list = False
            formatted_lines.append("")
            continue

        # Check for headings (e.g., "Target Market:")
        if line.endswith(":") and not line.startswith("-"):
            formatted_lines.append(f"### {line}")
            continue

        # Check for list items (e.g., "- Item" or "1. Item")
        if line.startswith("-") or line[0].isdigit() and line[1:3] in [". ", ".("]:
            in_list = True
            formatted_lines.append(line)
            continue

        # If not a heading or list item, treat as a paragraph
        if in_list:
            in_list = False
            formatted_lines.append("")
        formatted_lines.append(line)

    return "\n".join(formatted_lines)

def process_problem(problem, technique_name, model_name, role="", emotion=""):
    if not problem:
        return "Please enter a problem statement."
    
    technique = TECHNIQUE_CONFIGS.get(technique_name)
    if not technique:
        return f"Technique {technique_name} not found."
    
    llm_config = MODEL_CONFIGS.get(MODELS[model_name])
    if not llm_config:
        return f"Model {model_name} not found."
    
    try:
        # Handle techniques that require additional parameters
        kwargs = {"llm_config": llm_config}
        if technique_name == "RolePrompting":
            kwargs["role"] = role or "Expert"
        elif technique_name == "EmotionPrompting":
            kwargs["emotion"] = emotion or "thoughtful and methodical"
        elif technique_name == "Expert Chain-of-Thought":
            kwargs["role"] = role or "Expert"
        
        response = technique.execute(problem, **kwargs)
        # Format the response as Markdown
        markdown_response = format_as_markdown(response)
        return markdown_response
    except Exception as e:
        return f"**Error**: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Proctor AI Prompt Engineering App") as interface:
    gr.Markdown("# Proctor AI Prompt Engineering App")
    gr.Markdown("Enter a problem, select a technique and model, and get a response powered by OpenRouter.")
    
    problem_input = gr.Textbox(label="Problem Statement", placeholder="e.g., How to build a house for a family of 4?")
    technique_dropdown = gr.Dropdown(choices=list(TECHNIQUE_CONFIGS.keys()), label="Prompting Technique")
    model_dropdown = gr.Dropdown(choices=list(MODELS.keys()), label="Model")
    role_input = gr.Textbox(label="Role (for RolePrompting or Expert CoT)", placeholder="e.g., Expert House Builder", visible=False)
    emotion_input = gr.Textbox(label="Emotion (for EmotionPrompting)", placeholder="e.g., thoughtful and methodical", visible=False)
    output = gr.Markdown(label="Response")  # Changed to gr.Markdown for proper rendering
    
    submit_button = gr.Button("Generate Response")
    
    # Dynamic visibility for role and emotion inputs
    def update_inputs(technique):
        return {
            role_input: gr.update(visible=technique in ["RolePrompting", "Expert Chain-of-Thought"]),
            emotion_input: gr.update(visible=technique == "EmotionPrompting")
        }
    
    technique_dropdown.change(fn=update_inputs, inputs=technique_dropdown, outputs=[role_input, emotion_input])
    
    submit_button.click(
        fn=process_problem,
        inputs=[problem_input, technique_dropdown, model_dropdown, role_input, emotion_input],
        outputs=output
    )

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