import os
import json
import re
from huggingface_hub import InferenceClient
import gradio as gr
from pydantic import BaseModel, Field
from typing import Optional, Literal

class PromptInput(BaseModel):
    text: str = Field(..., description="The initial prompt text")
    meta_prompt_choice: Literal["star","done","physics","morphosis", "verse", "phor","bolism","math","arpe"] = Field(..., description="Choice of meta prompt strategy")

class RefinementOutput(BaseModel):
    query_analysis: Optional[str] = None
    initial_prompt_evaluation: Optional[str] = None
    refined_prompt: Optional[str] = None
    explanation_of_refinements: Optional[str] = None
    raw_content: Optional[str] = None

class PromptRefiner:
    def __init__(self, api_token: str):
        self.client = InferenceClient(token=api_token)

    def refine_prompt(self, prompt_input: PromptInput) -> RefinementOutput:
        if prompt_input.meta_prompt_choice == "morphosis":
            selected_meta_prompt = original_meta_prompt
        elif prompt_input.meta_prompt_choice == "verse":
            selected_meta_prompt = new_meta_prompt
        elif prompt_input.meta_prompt_choice == "physics":
            selected_meta_prompt = metaprompt1
        elif prompt_input.meta_prompt_choice == "bolism":
            selected_meta_prompt = loic_metaprompt
        elif prompt_input.meta_prompt_choice == "done":
            selected_meta_prompt = metadone
        elif prompt_input.meta_prompt_choice == "star":
            selected_meta_prompt = echo_prompt_refiner
        elif prompt_input.meta_prompt_choice == "math":
            selected_meta_prompt = math_meta_prompt  
        elif prompt_input.meta_prompt_choice == "arpe":
            selected_meta_prompt = autoregressive_metaprompt
        else:
            selected_meta_prompt = advanced_meta_prompt
    
        messages = [
            {"role": "system", "content": 'You are an expert at refining and extending prompts. Given a basic prompt, provide a more detailed.'},
            {"role": "user", "content": selected_meta_prompt.replace("[Insert initial prompt here]", prompt_input.text)}
        ]
        response = self.client.chat_completion(
            model=prompt_refiner_model,
            messages=messages,
            max_tokens=2000,
            temperature=0.8
        )
        response_content = response.choices[0].message.content.strip()
        try:
            json_match = re.search(r'<json>\s*(.*?)\s*</json>', response_content, re.DOTALL)
            if json_match:
                json_str = json_match.group(1)
                json_str = re.sub(r'\n\s*', ' ', json_str)
                json_str = json_str.replace('"', '\\"')
                json_output = json.loads(f'"{json_str}"')
                if isinstance(json_output, str):
                    json_output = json.loads(json_output)
                for key, value in json_output.items():
                    if isinstance(value, str):
                        json_output[key] = value.replace('\\"', '"')
                return RefinementOutput(**json_output, raw_content=response_content)
            else:
                raise ValueError("No JSON found in the response")
        except (json.JSONDecodeError, ValueError) as e:
            print(f"Error parsing JSON: {e}")
            print(f"Raw content: {response_content}")
            output = {}
            for key in ["initial_prompt_evaluation", "refined_prompt", "explanation_of_refinements"]:
                pattern = rf'"{key}":\s*"(.*?)"(?:,|\}})'
                match = re.search(pattern, response_content, re.DOTALL)
                if match:
                    output[key] = match.group(1).replace('\\n', '\n').replace('\\"', '"')
                else:
                    output[key] = ""
            return RefinementOutput(**output, raw_content=response_content)
                
    def apply_prompt(self, prompt: str, model: str) -> str:
        try:
            messages = [
                {"role": "system", "content": "You are a helpful assistant. Answer in stylized version with latex format or markdown if relevant. Separate your answer into logical sections using level 2 headers (##) for sections and bolding (**) for subsections.Incorporate a variety of lists, headers, and text to make the answer visually appealing"},
                {"role": "user", "content": prompt}
            ]

            response = self.client.chat_completion(
                model=model,
                messages=messages,
                max_tokens=2000,
                temperature=0.8
            )

            output = response.choices[0].message.content.strip()
            output = output.replace('\n\n', '\n').strip()
            return output
        except Exception as e:
            return f"Error: {str(e)}"

class GradioInterface:
    def __init__(self, prompt_refiner: PromptRefiner):
        self.prompt_refiner = prompt_refiner

        # Define custom CSS for containers
        custom_css = """
        .container {
            border: 2px solid #2196F3;
            border-radius: 10px;
            padding: 20px;
            margin: 15px;
            background: white;
            position: relative;
        }
        
        .container::before {
            position: absolute;
            top: -12px;
            left: 20px;
            background: white;
            padding: 0 10px;
            color: #2196F3;
            font-weight: bold;
            font-size: 1.2em;
        }

        /* Remove default Gradio styles */
        .no-background > div:first-child {
            border: none !important;
            background: transparent !important;
            box-shadow: none !important;
        }

        .title-container::before { content: ''; }
        .input-container::before { content: 'PROMPT REFINEMENT'; }
        .analysis-container::before { content: 'ANALYSIS & REFINEMENT'; }
        .model-container::before { content: 'MODEL APPLICATION'; }
        .results-container::before { content: 'RESULTS'; }
        .examples-container::before { content: 'EXAMPLES'; }

        /* Custom styling for radio buttons */
        .radio-group {
            display: flex;
            gap: 10px;
            margin: 10px 0;
        }
        """

        with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as self.interface:
            # Title Container
            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. Enter a main idea for a prompt, choose a meta prompt, and the model will attempt to generate an improved version.")

            # Input Container
            with gr.Column(elem_classes=["container"]):
                prompt_text = gr.Textbox(
                    label="Type the prompt (or let it empty to see metaprompt)",
                 #   elem_classes="no-background"
                )
                with gr.Accordion("Meta Prompt explanation", open=False):
                    gr.Markdown(explanation_markdown)
                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")

            # Analysis Container
            with gr.Column(elem_classes=["container", "analysis-container"]):
                gr.Markdown("### Initial prompt analysis")
                analysis_evaluation = gr.Markdown()
                gr.Markdown("### Refined Prompt")
                refined_prompt = gr.Textbox(
                    interactive=False,
                    elem_classes="no-background"
                )
                gr.Markdown("### Explanation of Refinements")
                explanation_of_refinements = gr.Markdown()
            
            with gr.Accordion("Full Response JSON", open=False, visible=False):
                full_response_json = gr.JSON()

            # Model Application Container
            with gr.Column(elem_classes=["container", "model-container"]):
                gr.Markdown("## See MetaPrompt Impact")            
                with gr.Row():
                    apply_model = gr.Dropdown(
                        [
                            "Qwen/Qwen2.5-72B-Instruct",
                            "meta-llama/Meta-Llama-3-70B-Instruct",
                            "meta-llama/Llama-3.1-8B-Instruct",
                            "NousResearch/Hermes-3-Llama-3.1-8B",
                            "HuggingFaceH4/zephyr-7b-alpha",
                            "meta-llama/Llama-2-7b-chat-hf",
                            "microsoft/Phi-3.5-mini-instruct"
                        ],
                        value="meta-llama/Meta-Llama-3-70B-Instruct",
                        label="Choose the Model to apply to the prompts (the one you will used)",
                        elem_classes="no-background"
                    )
                    apply_button = gr.Button("Apply MetaPrompt")

            # Results Container
            with gr.Column(elem_classes=["container", "results-container"]):
                with gr.Tabs():
                    with gr.TabItem("Original Prompt Output"):
                        original_output = gr.Markdown()
                    with gr.TabItem("Refined Prompt Output"):
                        refined_output = gr.Markdown()

            # Examples Container
            with gr.Column(elem_classes=["container", "examples-container"]):
                with gr.Accordion("Examples", open=True):
                    gr.Examples(
                        examples=[
                            ["Write a story on the end of prompt engineering replaced by an Ai specialized in refining prompts.", "star"],
                            ["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]
                    )

            # Connect the buttons to their functions
            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]
            )

    # Your existing methods remain the same
    def refine_prompt(self, prompt: str, meta_prompt_choice: str) -> tuple:
        input_data = PromptInput(text=prompt, meta_prompt_choice=meta_prompt_choice)
        result = self.prompt_refiner.refine_prompt(input_data)
        analysis_evaluation = f"\n\n{result.initial_prompt_evaluation}"
        return (
            analysis_evaluation,
            result.refined_prompt,
            result.explanation_of_refinements,
            result.dict()
        )

    def apply_prompts(self, original_prompt: str, refined_prompt: str, model: str):
        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

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

metaprompt_explanations = {
    "star": "Use ECHO when you need a comprehensive, multi-stage approach for complex prompts. It's ideal for tasks requiring in-depth analysis, exploration of multiple alternatives, and synthesis of ideas. Choose this over others when you have time for a thorough refinement process and need to consider various aspects of the prompt.",
    "done": "Opt for this when you want a structured approach with emphasis on role-playing and advanced techniques. It's particularly useful for tasks that benefit from diverse perspectives and complex reasoning. Prefer this over 'physics' when you need a more detailed, step-by-step refinement process.",
    "physics": "Select this when you need a balance between structure and advanced techniques, with a focus on role-playing. It's similar to 'done' but may be more suitable for scientific or technical prompts. Choose this over 'done' for a slightly less complex approach.",
    "morphosis": "Use this simplified approach for straightforward prompts or when time is limited. It focuses on essential improvements without complex techniques. Prefer this over other methods when you need quick, clear refinements without extensive analysis.",
    "verse": "Choose this method when you need to analyze and improve a prompt's strengths and weaknesses, with a focus on information flow. It's particularly useful for enhancing the logical structure of prompts. Use this over 'morphosis' when you need more depth but less complexity than 'star'.",
    "phor": "Employ this advanced approach when you need to combine multiple prompt engineering techniques. It's ideal for complex tasks requiring both clarity and sophisticated prompting methods. Select this over 'star' when you want a more flexible, technique-focused approach.",
    "bolism": "Utilize this method when working with autoregressive language models and when the task requires careful reasoning before conclusions. It's best for prompts that need detailed output formatting. Choose this over others when the prompt's structure and reasoning order are crucial."
}

explanation_markdown = "".join([f"- **{key}**: {value}\n" for key, value in metaprompt_explanations.items()])

# Main code to run the application
if __name__ == '__main__':
    meta_info=""
    api_token = os.getenv('HF_API_TOKEN')
    if not api_token:
        raise ValueError("HF_API_TOKEN not found in environment variables")

    metadone = os.getenv('metadone')
    prompt_refiner_model = os.getenv('prompt_refiner_model')
    echo_prompt_refiner = os.getenv('echo_prompt_refiner')
    metaprompt1 = os.getenv('metaprompt1')   
    loic_metaprompt = os.getenv('loic_metaprompt')    
    openai_metaprompt = os.getenv('openai_metaprompt')
    original_meta_prompt = os.getenv('original_meta_prompt')    
    new_meta_prompt = os.getenv('new_meta_prompt')   
    advanced_meta_prompt = os.getenv('advanced_meta_prompt')
    math_meta_prompt = os.getenv('metamath')
    autoregressive_metaprompt = os.getenv('autoregressive_metaprompt')


    prompt_refiner = PromptRefiner(api_token)
    gradio_interface = GradioInterface(prompt_refiner)
    gradio_interface.launch(share=True)