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

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

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}
"""

llm_prompt = PromptTemplate.from_template(llm_prompt_template)

gemini_embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")

llm = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.7, top_p=1)

questions_template = ChatPromptTemplate.from_template(template=llm_prompt_template)

questions_chain = LLMChain(llm=llm, prompt=questions_template)

questions_schema = ResponseSchema(
    name="questions",
    description="""Give the questions in json as an array""",
)

questions_schemas = [questions_schema]

questions_parser = StructuredOutputParser.from_response_schemas(questions_schemas)
format_questions_instructions = questions_parser.get_format_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"
    }
  ]
```
}"""