import gradio as gr
from llm_handler import generate_response

# Predefined system prompts
SYSTEM_PROMPTS = [
    "You are a professional coding assistant",
    "You are a professional teaching assistant",
    "You are professional email writer",
    "Other"
]

# Available models
MODELS = [
    "Claude Haiku",
    "DeepSeek",
    "Claude Premium",
    "GPT Pro"
]


def handle_system_prompt_change(choice):
    if choice == "Other":
        return gr.Textbox(visible=True)
    return gr.Textbox(visible=False)


def process_input(user_prompt, model_choice, system_prompt_choice, custom_system_prompt):
    # Use custom system prompt if "Other" is selected
    final_system_prompt = custom_system_prompt if system_prompt_choice == "Other" else system_prompt_choice

    if not user_prompt:
        return "Please enter a prompt"

    # Generate response using llmhandler
    response = generate_response(user_prompt, model_choice, final_system_prompt)
    return response


# Create Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# LLM Interface")

    with gr.Row():
        with gr.Column():
            user_prompt = gr.Textbox(
                label="Your Prompt",
                placeholder="Enter your prompt here...",
                lines=5
            )

            model_choice = gr.Dropdown(
                choices=MODELS,
                label="Select Model",
                value=MODELS[0]
            )

            system_prompt_choice = gr.Dropdown(
                choices=SYSTEM_PROMPTS,
                label="Select System Prompt",
                value=SYSTEM_PROMPTS[0]
            )

            custom_system_prompt = gr.Textbox(
                label="Custom System Prompt",
                placeholder="Enter your custom system prompt...",
                visible=False
            )

            submit_btn = gr.Button("Generate Response")

        with gr.Column():
            output = gr.Textbox(
                label="Response",
                lines=10,
                show_copy_button=True  # Adding copy button to the output
            )

    # Handle system prompt change
    system_prompt_choice.change(
        fn=handle_system_prompt_change,
        inputs=[system_prompt_choice],
        outputs=[custom_system_prompt]
    )

    # Handle submit button click
    submit_btn.click(
        fn=process_input,
        inputs=[
            user_prompt,
            model_choice,
            system_prompt_choice,
            custom_system_prompt
        ],
        outputs=[output]
    )

if __name__ == "__main__":
    demo.launch()