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, Tuple, Any
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")
    score: float = Field(..., description="Similarity score of the retrieved 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 vector_store


# 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
# Function to initialize prompt template
def initialize_prompt_template(parser):
    prompt_template = """
    You are given a **query** and relevant **context** below. Based on both, generate questions.

    **Query**: {query}

    **Context**: {context}

    **Question Type Instructions:**
    - If a **specific question type** is provided, generate questions **only in that type**: {question_type}.
    - Do **not** generate other question types if a type is specified.
    - 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 the generated questions align with the query and the retrieved context.**

    The response should be structured in this format:
    
    - 'question': The question prompt.
    - 'type': The type of question (should match the requested type, unless 'general').
    - 'options': For 'MCQ', a list of answer options (omit for other types).

    **Strictly follow this structured format**:
    {format_instructions}
    """
    return prompt_template



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

# Function for retrieving with score
def retrieve_and_format_results(vector_store: Any, query: str, k: int = 5, filter: dict = {}) -> RetrievedDocsSchema:
    """
    Retrieves documents using similarity search and formats them into the Pydantic schema.

    Args:
        vector_store (Any): The vector store used for retrieval.
        query (str): The search query.
        k (int): Number of documents to retrieve.
        filter (dict): Optional filter for the search.

    Returns:
        RetrievedDocsSchema: A structured schema containing documents and metadata.
    """
    # Retrieve documents with similarity scores
    retrieved_docs = vector_store.similarity_search_with_score(query, k=k, filter=filter)

    # Convert retrieved documents into the Pydantic schema
    documents_list = [
        DocumentSchema(
            metadata=MetadataSchema(
                page=doc.metadata.get("page", 0),
                page_label=doc.metadata.get("page_label", ""),
                total_pages=doc.metadata.get("total_pages", 0),
                source=doc.metadata.get("source", ""),
                score=score  # Assign the similarity score
            ),
            page_content=doc.page_content
        )
        for doc, score in retrieved_docs
    ]

    return RetrievedDocsSchema(documents=documents_list)

# llm Generation function
def generate_questions_from_context(query: str, vector_store: Any, llm_structured, prompt_template: str,
    parser: PydanticOutputParser, chat_history: List[dict], question_type: str = "general") -> Tuple[Any, RetrievedDocsSchema]:
    """
    Generates questions based on retrieved document context using an LLM.

    Args:
        query (str): The search query.
        vector_store (Any): The vector store used for retrieval.
        llm_structured: The structured LLM output function.
        prompt_template (str): The prompt template for question generation.
        parser (PydanticOutputParser): The Pydantic output parser.
        chat_history (List[dict]): A list to store conversation history.
        question_type (str): The type of question (default is "general").

    Returns:
        Tuple[Any, RetrievedDocsSchema]: A tuple containing the LLM-generated response and retrieved document schema.
    """
    # Append user query to chat history
    chat_history.append({"role": "user", "content": query})

    # Retrieve and format results using structured schema
    retrieved_docs_schema = retrieve_and_format_results(vector_store, query, k=5, filter={})

    # Extract only the page_content from the retrieved documents
    retrieved_docs = [doc.page_content for doc in retrieved_docs_schema.documents]

    # Combine retrieved documents into a single context string
    context = " ".join(retrieved_docs)

    # Initialize the prompt with query, context, and format instructions
    prompt = PromptTemplate(
        template=prompt_template,
        input_variables=["query", "context", "question_type"],
        partial_variables={"format_instructions": parser.get_format_instructions()}
    )

    # Format the prompt with input variables
    formatted_prompt = prompt.format(
        query=query,
        context=context,
        question_type=question_type
    )

    # Generate response using the LLM
    chain = prompt | llm_structured
    response = chain.invoke({
        "query": query,
        "context": context,
        "question_type": question_type,
        "format_instructions": parser.get_format_instructions()
    })


    # Append assistant response to chat history
    chat_history.append({"role": "assistant", "content": str(response)})

    return response, retrieved_docs_schema, formatted_prompt


# Streamlit interface
import streamlit as st

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

    # Initialize chat history in session state
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []

    # Sidebar for 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']}")

    # Dropdown for Question Type Selection
    question_type = st.selectbox("Select Question Type", ["general", "MCQ", "fill_missing", "short_answer"])

    # Input for Query
    query = st.chat_input("Enter your query: ")

    if query and question_type:
        with st.spinner('Generating questions...'):
            st.write(f"**Your query:** {query}")

            # Call the updated function that now returns the generated prompt as well
            response, retrieved_docs_schema, generated_prompt = generate_questions_from_context(
                query, vector_store, llm_structured, prompt_template, parser, st.session_state.chat_history, question_type
            )

            # Display Generated Prompt
            st.subheader("Generated Prompt")
            # st.code(generated_prompt, language="plaintext")
            st.write(f"**Final prompt is:** {generated_prompt}")

            # Display Generated Questions
            st.subheader("Generated Questions")
            st.write(response)  # Displaying as structured JSON for clarity

            # Display Retrieved Documents with Scores
            st.subheader("Retrieved Documents")
            for doc in retrieved_docs_schema.documents:
                with st.expander(f"Source: {doc.metadata.source}, Page {doc.metadata.page}/{doc.metadata.total_pages} (Score: {doc.metadata.score:.4f})"):
                    st.text_area("Content:", doc.page_content, height=150)

if __name__ == "__main__":
    main()