import streamlit as st
from langchain_pinecone import PineconeVectorStore
from langchain_openai import OpenAI, ChatOpenAI
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.prompts import PromptTemplate
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.output_parsers import PydanticOutputParser
from dotenv import load_dotenv
import os
from pydantic import BaseModel, Field
from typing import List, Union
from typing_extensions import Literal


# Pydantic Schema
class Question(BaseModel):
    question: str = Field(..., description="The question prompt that the user needs to answer.")
    type: Literal['fill_missing', 'MCQ', 'short_answer'] = Field(..., description="The type of question: fill_missing, MCQ, or short_answer.")
    options: Union[List[str], None] = Field(None, description="The options for the question, used only for MCQ type.")

class Questions(BaseModel):
    no_of_questions: int = Field(..., description="The total number of questions generated.")
    questions: List[Question] = Field(..., description="A list of Question objects.")


# Function to download Hugging Face embeddings
def download_hugging_face_embeddings():
    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
    return embeddings


# Load environment variables
def load_env_variables():
    load_dotenv()
    PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
    OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
    GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")

    os.environ["PINECONE_API_KEY"] = PINECONE_API_KEY
    os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
    os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY


# Function to initialize Pinecone vector store and retriever
def initialize_vector_store(embeddings):
    index_name = "yolotest"
    vector_store = PineconeVectorStore.from_existing_index(index_name=index_name, embedding=embeddings)
    retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
    return retriever


# Function to initialize LLM
def initialize_llm():
    llm = ChatOpenAI(api_key=os.environ.get("OPENAI_API_KEY"), temperature=0, model='gpt-3.5-turbo-0125')
    llm_structured = llm.with_structured_output(Questions)
    return llm_structured


# Function to initialize prompt template
def initialize_prompt_template(parser):
    prompt_template = """
    You are given some context below. Based on the context, generate questions with one of the following types: ['fill_missing', 'MCQ', 'short_answer'].

    Context: {context}

    Make sure to generate a variety of questions. The types should be distributed among 'fill_missing', 'MCQ', and 'short_answer'. 

    For MCQs, provide at least 3 options. For 'fill_missing', make sure to leave a gap that can be filled. For 'short_answer', make sure the answer is clear from the context.

    Return a list of questions, with each question having the following structure:
    - 'question': The question prompt.
    - 'type': The type of question: 'fill_missing', 'MCQ', or 'short_answer'.
    - 'options': For 'MCQ', a list of options. For other types, this field should be omitted.

    The response should be with a format that matches the following structure:
    {format_instructions}
    """
    return prompt_template


# Initialize components before user query
embeddings = download_hugging_face_embeddings()
load_env_variables()
retriever = initialize_vector_store(embeddings)
llm_structured = initialize_llm()
parser = PydanticOutputParser(pydantic_object=Questions)
prompt_template = initialize_prompt_template(parser)


# Function to generate questions from the context
def generate_questions_from_context(query: str, retriever, llm_structured, prompt_template, parser):
    # Retrieve relevant documents from the vector store using the retriever
    retrieved_docs = retriever.invoke(query)
    retrieved_docs = [rec.page_content for rec in retrieved_docs]

    # Combine the retrieved documents into a single context string
    context = " ".join([doc for doc in retrieved_docs])

    # Initialize prompt template
    prompt = PromptTemplate(template=prompt_template, input_variables=["context"],
                            partial_variables={"format_instructions": parser.get_format_instructions()})

    # Create chain and generate response
    chain = prompt | llm_structured
    response = chain.invoke({
        "context": context, "format_instructions": parser.get_format_instructions()
    })
    return response


# Streamlit Interface
def main():
    st.title("A Simple RAG App to Generate Questions in Specific Formats")

    # Display user query input
    query = st.chat_input("Say something: ")
    if query:
        with st.spinner('Generating questions...'):
            st.write(f"Your query: {query}")  # Display user query
            
            # Generate questions
            result = generate_questions_from_context(query, retriever, llm_structured, prompt_template, parser)
            st.write(result)
            # # Display the result in a more readable format
            # st.subheader("Generated Questions")
            # for i, question in enumerate(result.questions):
            #     st.markdown(f"### Question {i + 1}: {question.question}")
            #     st.markdown(f"**Type**: {question.type}")
            #     if question.type == "MCQ" and question.options:
            #         st.markdown("**Options**:")
            #         for option in question.options:
            #             st.markdown(f"- {option}")
            #     st.markdown("---")


# Run the app
if __name__ == "__main__":
    main()