import gradio as gr
import google.generativeai as genai
import time
import yaml
import os
from typing import Any, Generator, Tuple
from pydantic import BaseModel, ValidationError

# Wczytywanie klucza API z zmiennych środowiskowych
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")

# Sprawdzenie, czy klucz API został znaleziony
if not GEMINI_API_KEY:
    raise ValueError("Nie znaleziono klucza API Gemini. Ustaw zmienną środowiskową GEMINI_API_KEY.")

genai.configure(api_key=GEMINI_API_KEY)

generation_config = {
    "temperature": 1,
    "top_p": 0.95,
    "top_k": 64,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}

system_message = """Jesteś "Nauczycielem Promptów", asystentem, który pomaga w pisaniu lepszych promptów. Twoim zadaniem jest ulepszanie promptów, aby były bardziej precyzyjne i dawały lepsze wyniki. 
Kiedy otrzymasz prośbę o ulepszenie promptu, odpowiedz tylko ulepszoną wersją promptu. 
Kiedy otrzymasz prośbę o wyjaśnienie zmian, opisz je w prosty i zrozumiały sposób."""

inital_usr_text = "Chcę nauczyć się pisać lepsze prompty! Pomożesz mi? 🤔📚"

initial_bot_text = """Witaj! 👋 Jestem Nauczycielem Promptów -  pomogę Ci ulepszyć Twoje prompty! 🤖🎓
Zobacz, jak to działa:
1. Wybierz, co chcesz zmienić w swoim prompcie. 
2. Wpisz prompt - co chcesz, żeby komputer zrobił? 
3. Kliknij "✨Ulepsz prompt!", a ja postaram się go ulepszyć! 
4. Kliknij "Wyjaśnij zmiany 💡", jeśli chcesz zrozumieć, co zmieniłem.
Możesz też zobaczyć przykłady na dole strony - kliknij na któryś, żeby go wybrać. 🖱️👇"""


class Metaprompt(BaseModel):
    explanation: str
    example_prompt: str
    example_prompt_explanation: str
    name: str
    template: str

    def __str__(self):
        return f"✨**{self.name}**\n\n📝 *{self.explanation}*\n\n📚 **Przykładowy Prompt:** {self.example_prompt}\n\n📖 **Wyjaśnienie Przykładowego Promptu:** {self.example_prompt_explanation}"


class MetapromptLibrary(BaseModel):
    Metaprompts: list[Metaprompt]


def read_and_validate(file_path: str = "metaprompts.yml"):
    try:
        with open(file_path, "r") as file:
            data = yaml.safe_load(file)
            validated_data = MetapromptLibrary(**data)
            return validated_data
    except FileNotFoundError:
        raise FileNotFoundError(
            "Nie znaleziono pliku metaprompts.yml. Upewnij się, że plik znajduje się w tym samym folderze co skrypt.")
    except yaml.YAMLError as e:
        raise yaml.YAMLError(f"Błąd parsowania pliku YAML: {e}")
    except ValidationError as e:
        raise ValidationError(f"Błąd walidacji danych: {e}")


metaprompts = read_and_validate().Metaprompts
metaprompts_dict = {mp.name: mp for mp in metaprompts}


def get_llm():
    return genai.GenerativeModel(
        model_name="gemini-1.5-pro-exp-0827",
        generation_config=generation_config,
    )


def explain_metaprompt(explanation_history, metaprompt):
    explanation_history += [
        [f"❓Jak **{metaprompt.capitalize()}** ulepsza mój prompt? 💡🚀", ""]
    ]
    answer = f"""{metaprompts_dict[metaprompt]}"""

    for character in answer:
        explanation_history[-1][1] += character
        time.sleep(0.001)
        yield explanation_history


def update_widgets(metaprompt, feedback):
    button_variant = "primary" if metaprompt else "secondary"
    feedback_visibility = True if metaprompt == "Dodaj uwagi" else False
    return [
        gr.Button(variant=button_variant),
        gr.Textbox(
            visible=feedback_visibility, value=feedback if feedback_visibility else ""
        ),
    ]


def explain_improvement(
        prompt, metaprompt, improved_prompt, prompt_teacher
):
    llm = get_llm()

    chat_session = llm.start_chat()

    response = chat_session.send_message(
        f"""Następujący prompt: 
---
{prompt} 
---
został ulepszony za pomocą metapromptu:
---
{metaprompt}
---
W rezultacie powstała następująca ulepszona wersja:
---
{improved_prompt}
---
Zwięźle wyjaśnij ulepszenie."""
    )

    prompt_teacher += [
        [
            "❓Czy możesz proszę **wyjaśnić** **ulepszony prompt**? 📝✨",
            "**Wyjaśnienie**:\n\n",
        ]
    ]

    prompt_teacher[-1][1] += response.text
    yield prompt_teacher


def improve_prompt(
        prompt: str,
        metaprompt: str,
        feedback: str | None,
        explanation_history,
) -> Generator[Tuple[str, str], Any, Any]:
    metaprompt_template = metaprompts_dict[metaprompt].template

    llm = get_llm()
    chat_session = llm.start_chat()

    chat_session.send_message(system_message)

    explanation_history += [
        [
            f"""
            **📝 Proszę ulepszyć następujący Prompt:**
            ```
            {prompt}
            ```
            """,
            "**Ulepszony Prompt**:\n\n",
        ]
    ]

    improved_prompt = ""

    input_message = {"prompt": prompt, "feedback": feedback} if feedback else {"prompt": prompt}
    input_str = ""
    for key, value in input_message.items():
        input_str += f"{key}: {value}\n"

    response = chat_session.send_message(metaprompt_template + "\n" + input_str)

    explanation_history[-1][1] += response.text
    improved_prompt += response.text
    yield improved_prompt, explanation_history


def robustly_improve_prompt(*args, **kwargs):
    history = args[3]
    user_txt = "Ups, wystąpił błąd!💥 Co powinienem zrobić?"
    try:
        yield from improve_prompt(*args, **kwargs)
    except Exception as e:
        msg = "\n\nSpróbuj jeszcze raz! Upewnij się, że wszystko jest dobrze wpisane. 🤔"
        try:
            error_msg = e.body["message"]
            error_msg += msg
            history += [[user_txt, error_msg]]
            gr.Warning(error_msg)
            yield "", history
        except Exception as ee:
            error_msg = str(e) + msg
            history += [[user_txt, error_msg]]
            gr.Warning(error_msg)
            yield "", history


with gr.Blocks(title="Nauczyciel Promptów", theme=gr.themes.Soft()) as gradio_app:
    # Zastosowanie gr.HTML do wyświetlenia HTML
    gr.HTML(
        "<center><h1>🤖 <font size = '7'>Nauczyciel Promptów</font> - Pomocnik w pisaniu lepszych promptów! 📝✨ </h1></center>")
    with gr.Accordion("ℹ️ Co to są prompty?", open=True):
        gr.Markdown(
            """
        1. **Co to jest prompt?**
        Prompt to jakby instrukcja, którą dajemy komputerowi. To tak, jakbyś mówił przyjacielowi, co ma zrobić. Na przykład, jeśli chcesz, żeby przyjaciel narysował kota, mówisz mu: „Narysuj kota z dużymi oczami i puszystym ogonem.” To właśnie jest twój prompt!

        2. **Jak działa prompt?**
        Kiedy wpisujesz prompt w komputerze, on czyta go i stara się zrobić to, co napisałeś. Im bardziej szczegółowy i jasny jest twój prompt, tym lepszy obrazek lub odpowiedź dostaniesz.

        3. **Dlaczego dobry prompt jest ważny?**
        Jeśli powiesz komputerowi coś ogólnego, na przykład „narysuj coś”, to może narysować cokolwiek, co niekoniecznie będzie tym, co chcesz. Ale jeśli powiesz: „narysuj piękny zachód słońca z palmami i plażą”, komputer będzie miał lepszy pomysł na to, co powinien stworzyć.

        4. **Przykłady dobrych promptów:**

        „Napisz krótką bajkę o małym smoku, który szuka przyjaciół.”
        „Stwórz obrazek, na którym są dinozaury bawiące się w parku.”
        5. **Ćwiczenie!**
        Spróbuj stworzyć własny prompt! Co byś chciał, żeby komputer narysował lub napisał? Pamiętaj, im więcej szczegółów dodasz, tym lepszy efekt dostaniesz!
        """
        )
    with gr.Row():
        with gr.Column(scale=2):
            prompt_teacher = gr.Chatbot(
                height=580,
                label="Nauczyciel Promptów",
                show_copy_button=True,
                value=[[inital_usr_text, initial_bot_text]],
                avatar_images=("thinking.svg", "robot.svg"), # Ustawienie avatarów
            )
            prompt = gr.Textbox(
                label="Prompt",
                interactive=True,
                placeholder="Wpisz swój prompt - co chcesz, żeby komputer zrobił?",
                value="Napisz krótkie streszczenie bitwy pod Grunwaldem.",
                show_copy_button=True,
            )
            with gr.Row():
                explain_btn = gr.Button(
                    "Wyjaśnij zmiany 💡",
                    variant="primary",
                    visible=False,
                )
                replace_btn = gr.Button(
                    "Zatwierdź zmiany 👍",
                    variant="primary",
                    visible=False,
                )
            with gr.Row():
                improve_btn = gr.Button("✨Ulepsz prompt!", variant="primary")
        with gr.Column(scale=1):
            metaprompt = gr.Radio(
                label="Ulepszenia",
                info="Wybierz, jak chcesz ulepszyć swój prompt.",
                value="Dodaj szczegóły",
                choices=[mp.name.replace("_", " ").capitalize() for mp in metaprompts],
            )
            feedback = gr.Textbox(
                label="Twoje uwagi",
                info="Napisz, co chcesz zmienić w prompcie.",
                visible=False,
            )

    improved_prompt = gr.Textbox(label="Ulepszony Prompt", visible=False)
    examples = gr.Examples(
        examples=[[mp.name, mp.example_prompt] for mp in metaprompts],
        examples_per_page=100,
        inputs=[metaprompt, prompt],
        label="Przykłady",
    )

    metaprompt.change(
        fn=update_widgets,
        inputs=[metaprompt, feedback],
        outputs=[improve_btn, feedback],
    ).success(
        lambda: [gr.Button(visible=False), gr.Button(visible=False)],
        None,
        [replace_btn, explain_btn],
    ).success(
        fn=explain_metaprompt,
        inputs=[prompt_teacher, metaprompt],
        outputs=[prompt_teacher],
    )
    improve_btn.click(
        fn=robustly_improve_prompt,
        inputs=[
            prompt,
            metaprompt,
            feedback,
            prompt_teacher,
        ],
        outputs=[improved_prompt, prompt_teacher],
    ).success(
        lambda: [gr.Button(visible=True), gr.Button(visible=True)],
        None,
        [replace_btn, explain_btn],
    )

    explain_btn.click(lambda: gr.Button(visible=False), None, explain_btn).success(
        explain_improvement,
        [
            prompt,
            metaprompt,
            improved_prompt,
            prompt_teacher,
        ],
        prompt_teacher,
    )

    replace_btn.click(lambda x: x, improved_prompt, prompt).success(
        lambda: [gr.Button(visible=False), gr.Button(visible=False)],
        None,
        [replace_btn, explain_btn],
    )

if __name__ == "__main__":
    gradio_app.queue(default_concurrency_limit=10).launch(favicon_path="robot.svg")