import os
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain import PromptTemplate, LLMChain
from langchain.output_parsers import ResponseSchema, StructuredOutputParser
from langchain.prompts import ChatPromptTemplate


class Gemini:

    _llm_prompt_template = """
    Olá, sou uma IA treinada para gerar conteúdo educacional. Por favor, gere cinco questões de múltipla escolha sobre o seguinte tema:
    Instruções para cada questão:
    - Crie uma questão clara e relevante para o tema.
    - Forneça cinco opções de resposta, rotuladas de A) a E).
    - Apenas uma das opções de resposta deve ser correta.
    - Indique a resposta correta ao final de cada questão.

    Exemplo de uma questão:
    Tema: Fotossíntese

    Questão:
    Qual é o pigmento primário responsável pela fotossíntese nas plantas?

    Opções de Resposta:
    A) Clorofila
    B) Hemoglobina
    C) Mioglobina
    D) Citocromo
    E) Queratina

    Resposta Correta:
    A) Clorofila

    Context: {context}
    Question: {question}
    Answer:


    {format_questions_instructions}
    """

    _format_questions_instructions = """
    The output should be a markdown code snippet formatted in the following schema, including the leading and trailing "```json" and "```":
    ```json
    {
      "questions": [
        {
          question: "Qual é o pigmento primário responsável pela fotossíntese nas plantas?",
          options: ["A) Clorofila",
                    "B) Hemoglobina",
                    "C) Mioglobina",
                    "D) Citocromo",
                    "E) Queratina"],
          answer: "A"
        }
      ]
    ```
    }"""

    def __init__(self):

        if "GOOGLE_API_KEY" not in os.environ:
            raise ValueError("GOOGLE_API_KEY environment variable is not set")

        self.llm_prompt = PromptTemplate.from_template(self._llm_prompt_template)

        self.embeddings_model = "models/embedding-001"
        self.model = "gemini-pro"

        self.embeddings = GoogleGenerativeAIEmbeddings(model=self.embeddings_model)
        self.llm = ChatGoogleGenerativeAI(model=self.model, temperature=0.7, top_p=1)

        self.template = ChatPromptTemplate.from_template(
            template=self._llm_prompt_template
        )

        self.chain = LLMChain(llm=self.llm, prompt=self.template)

        self.schemas = [
            ResponseSchema(
                name="questions",
                description="""Give the questions in json as an array""",
            )
        ]

        self.parser = StructuredOutputParser.from_response_schemas(self.schemas)