import gradio as gr
from prompt_refiner import PromptRefiner
from variables import models, explanation_markdown, custom_css

class GradioInterface:
    def __init__(self, prompt_refiner: PromptRefiner, custom_css):
        self.prompt_refiner = prompt_refiner
        with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as self.interface:
            with gr.Column(elem_classes=["container", "title-container"]):
                gr.Markdown("# PROMPT++")
                gr.Markdown("### Automating Prompt Engineering by Refining your Prompts")
                gr.Markdown("Learn how to generate an improved version of your prompts.")
                gr.HTML(
                "<p style='text-align: center; color:orange;'>⚠ This space is in progress, and we're actively working on it, so you might find some bugs! Please report any issues you have in the Community tab to help us make it better for all.</p>"
                )
        
            with gr.Column(elem_classes=["container", "input-container"]):
                prompt_text = gr.Textbox(
                    label="Type your prompt (or let it empty to see metaprompt)",
                    lines=5
                )
                meta_prompt_choice = gr.Radio(
                    ["star","done","physics","morphosis", "verse", "phor","bolism","math","arpe"],
                    label="Choose Meta Prompt",
                    value="star",
                    elem_classes=["no-background", "radio-group"]
                )
                refine_button = gr.Button("Refine Prompt")  
                
                with gr.Row(elem_classes=["container2"]):
                    with gr.Accordion("Examples", open=False):
                        gr.Examples(
                            examples=[
                                ["Write a story on the end of prompt engineering replaced by an Ai specialized in refining prompts.", "done"],
                                ["Tell me about that guy who invented the light bulb", "physics"],
                                ["Explain the universe.", "star"],
                                ["What's the population of New York City and how tall is the Empire State Building and who was the first mayor?", "morphosis"],
                                ["List American presidents.", "verse"],                        
                                ["Explain why the experiment failed.", "morphosis"],
                                ["Is nuclear energy good?", "verse"],
                                ["How does a computer work?", "phor"],
                                ["How to make money fast?", "done"],
                                ["how can you prove IT0's lemma in stochastic calculus ?", "arpe"],                    
                            ],
                            inputs=[prompt_text, meta_prompt_choice]
                        )
                       
                    with gr.Accordion("Meta Prompt explanation", open=False):
                        gr.Markdown(explanation_markdown)
                
            with gr.Column(elem_classes=["container", "analysis-container"]):
                gr.Markdown(' ')
                gr.Markdown("### Initial prompt analysis")
                analysis_evaluation = gr.Markdown()
                gr.Markdown("### Refined Prompt")
                refined_prompt = gr.Textbox(
                    label="Refined Prompt",
                    interactive=True,
                    show_label=True,
                    show_copy_button=True,
                )
                gr.Markdown("### Explanation of Refinements")
                explanation_of_refinements = gr.Markdown()
            
            with gr.Column(elem_classes=["container", "model-container"]):
                with gr.Row():
                    apply_model = gr.Dropdown(models,
                        value="meta-llama/Llama-3.1-8B-Instruct",
                        label="Choose the Model",
                        container=False,
                        scale=1,
                        min_width=300
                    )
                    apply_button = gr.Button("Apply MetaPrompt")

                gr.Markdown("### Prompts on choosen model")
                with gr.Tabs():
                    with gr.TabItem("Original Prompt Output"):
                        original_output = gr.Markdown()
                    with gr.TabItem("Refined Prompt Output"):
                        refined_output = gr.Markdown()
            with gr.Accordion("Full Response JSON", open=False, visible=True):
                full_response_json = gr.JSON()
                
            refine_button.click(
                fn=self.refine_prompt,
                inputs=[prompt_text, meta_prompt_choice],
                outputs=[analysis_evaluation, refined_prompt, explanation_of_refinements, full_response_json]
            )

            apply_button.click(
                fn=self.apply_prompts,
                inputs=[prompt_text, refined_prompt, apply_model],
                outputs=[original_output, refined_output],
                api_name="apply_prompts"
            )

    def refine_prompt(self, prompt: str, meta_prompt_choice: str) -> tuple:
        initial_prompt_evaluation, refined_prompt, explanation_refinements, full_response = self.prompt_refiner.refine_prompt(prompt, meta_prompt_choice)
        analysis_evaluation = f"\n\n{initial_prompt_evaluation}"
        return (
            analysis_evaluation,
            refined_prompt,
            explanation_refinements,
            full_response
        )

    def apply_prompts(self, original_prompt: str, refined_prompt: str, model: str):
        try:
            original_output = self.prompt_refiner.apply_prompt(original_prompt, model)
            refined_output = self.prompt_refiner.apply_prompt(refined_prompt, model)
            return original_output, refined_output
        except Exception as e:
            return f"Error: {str(e)}", f"Error: {str(e)}"

    def launch(self, share=False):
        self.interface.launch(share=share)