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

# Generation 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.")

# Retrieval schema
class MetadataSchema(BaseModel):
    page: float = Field(..., description="Page number of the document")
    page_label: str = Field(..., description="Page label of the document")
    total_pages: float = Field(..., description="Total pages in the document")
    source: str = Field(..., description="Source file path of the document")

class DocumentSchema(BaseModel):
    metadata: MetadataSchema = Field(..., description="Filtered metadata of the document")
    page_content: str = Field(..., description="Content of the document page")

class RetrievedDocsSchema(BaseModel):
    documents: List[DocumentSchema]


# 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.
    
    Context: {context}
    
    - If a question type is specified, generate questions **only in that type**: {question_type}.
    - If a question type is specified, Do not generate any other question type expect for the type mentioned
    - If 'MCQ', provide at least 3 options per question.
    - If 'fill_missing', leave a blank space for the missing word.
    - If 'short_answer', ensure the answer is clear from the context.
    - If no type is specified (or 'general' is selected), generate a variety of question types.

    **Ensure that the generated questions match the requested type.**
    
    The response should follow this structure:
    
    - 'question': The question prompt.
    - 'type': The type of question (should match the requested type, unless 'general').
    - 'options': For 'MCQ', a list of answer options (otherwise, omit this field).
    
    The response should strictly match this format:
    {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)


def retrieve_documents(retriever, query: str, k: int = 5) -> RetrievedDocsSchema:
    retrieved_docs = retriever.invoke(query)
    
    extracted_docs = [
        DocumentSchema(
            metadata=MetadataSchema(
                page=doc.metadata.get("page", 0.0),
                page_label=doc.metadata.get("page_label", ""),
                total_pages=doc.metadata.get("total_pages", 0.0),
                source=doc.metadata.get("source", "")
            ),
            page_content=doc.page_content
        )
        for doc in retrieved_docs
    ]
    
    return RetrievedDocsSchema(documents=extracted_docs)


def generate_questions_from_context(query: str, retriever, llm_structured, prompt_template, parser, chat_history, question_type=None):
    if question_type is None:
        question_type = "general"

    chat_history.append({"role": "user", "content": query})
    retrieved_docs_schema = retrieve_documents(retriever, query)
    retrieved_docs = [doc.page_content for doc in retrieved_docs_schema.documents]
    context = " ".join(retrieved_docs)
    
    prompt = PromptTemplate(
        template=prompt_template,
        input_variables=["context", "question_type"],
        partial_variables={"format_instructions": parser.get_format_instructions()}
    )
    
    chain = prompt | llm_structured
    response = chain.invoke({
        "context": context,
        "question_type": question_type,
        "format_instructions": parser.get_format_instructions()
    })
    
    chat_history.append({"role": "assistant", "content": str(response)})
    return response, retrieved_docs_schema

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

    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []

    with st.sidebar:
        st.subheader("Chat History")
        with st.expander("Show/Hide Chat History", expanded=False):
            for message in st.session_state.chat_history:
                st.markdown(f"**{message['role'].capitalize()}**: {message['content']}")

    question_type = st.selectbox("Select Question Type", ["general", "MCQ", "fill_missing", "short_answer"])
    query = st.chat_input("Enter your query: ")

    if query and question_type:
        with st.spinner('Generating questions...'):
            st.write(f"**Your query:** {query}")
            response, retrieved_docs_schema = generate_questions_from_context(
                query, retriever, llm_structured, prompt_template, parser, st.session_state.chat_history, question_type
            )
            
            st.subheader("Generated Questions")
            st.write(response)
            # st.write(question_type)
            st.subheader("Retrieved Documents")
            for doc in retrieved_docs_schema.documents:
                st.markdown(f"**Source:** {doc.metadata.source}, Page {doc.metadata.page}/{doc.metadata.total_pages}")
                st.text_area("Content:", doc.page_content, height=100)

            # st.subheader("Chat History")
            # for message in st.session_state.chat_history:
            #     st.markdown(f"**{message['role'].capitalize()}**: {message['content']}")

if __name__ == "__main__":
    main()