#!/usr/bin/env python3
"""
question_and_answer: Tool for performing Q&A on PDF documents using retrieval augmented generation.

This module provides functionality to extract text from PDF binary data, split it into 
chunks, retrieve relevant segments via a vector store, and generate an answer to a 
user-provided question using a language model chain.
"""

import io
import logging
from typing import Annotated, Dict, Any, List

from PyPDF2 import PdfReader
from pydantic import BaseModel, Field
import hydra

from langchain.chains.question_answering import load_qa_chain
from langchain.docstore.document import Document
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.vectorstores import Annoy
from langchain_openai import OpenAIEmbeddings
from langchain_core.language_models.chat_models import BaseChatModel

from langchain_core.messages import ToolMessage
from langchain_core.tools import tool
from langchain_core.tools.base import InjectedToolCallId
from langgraph.types import Command
from langgraph.prebuilt import InjectedState

# Set up logging.
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# Load configuration using Hydra.
with hydra.initialize(version_base=None, config_path="../../configs"):
    cfg = hydra.compose(config_name="config", overrides=["tools/question_and_answer=default"])
    cfg = cfg.tools.question_and_answer
    logger.info("Loaded Question and Answer tool configuration.")

class QuestionAndAnswerInput(BaseModel):
    """
    Input schema for the PDF Question and Answer tool.

    Attributes:
        question (str): The question to ask regarding the PDF content.
        tool_call_id (str): Unique identifier for the tool call, injected automatically.
    """
    question: str = Field(
        description="The question to ask regarding the PDF content."
    )
    tool_call_id: Annotated[str, InjectedToolCallId]
    state: Annotated[dict, InjectedState]

def extract_text_from_pdf_data(pdf_bytes: bytes) -> str:
    """
    Extract text content from PDF binary data.

    This function uses PyPDF2 to read the provided PDF bytes and concatenates the text
    extracted from each page.

    Args:
        pdf_bytes (bytes): The binary data of the PDF document.

    Returns:
        str: The complete text extracted from the PDF.
    """
    reader = PdfReader(io.BytesIO(pdf_bytes))
    text = ""
    for page in reader.pages:
        page_text = page.extract_text() or ""
        text += page_text
    return text

def generate_answer(question: str, pdf_bytes: bytes, llm_model: BaseChatModel) -> Dict[str, Any]:
    """
    Generate an answer for a question using retrieval augmented generation on PDF content.

    This function extracts text from the PDF data, splits the text into manageable chunks,
    performs a similarity search to retrieve the most relevant segments, and then uses a
    question-answering chain (built using the provided llm_model) to generate an answer.

    Args:
        question (str): The question to be answered.
        pdf_bytes (bytes): The binary content of the PDF document.
        llm_model (BaseChatModel): The language model instance to use for answering.

    Returns:
        Dict[str, Any]: A dictionary containing the answer generated by the language model.
    """
    text = extract_text_from_pdf_data(pdf_bytes)
    logger.info("Extracted text from PDF.")
    text_splitter = CharacterTextSplitter(
        separator="\n",
        chunk_size=cfg.chunk_size,
        chunk_overlap=cfg.chunk_overlap
    )
    chunks = text_splitter.split_text(text)
    documents: List[Document] = [Document(page_content=chunk) for chunk in chunks]
    logger.info("Split PDF text into %d chunks.", len(documents))

    embeddings = OpenAIEmbeddings(openai_api_key=cfg.openai_api_key)
    vector_store = Annoy.from_documents(documents, embeddings)
    search_results = vector_store.similarity_search(
        question,
        k=cfg.num_retrievals
    )
    logger.info("Retrieved %d relevant document chunks.", len(search_results))
    # Use the provided llm_model to build the QA chain.
    qa_chain = load_qa_chain(llm_model, chain_type=cfg.qa_chain_type)
    answer = qa_chain.invoke(
        input={"input_documents": search_results, "question": question}
    )
    return answer

@tool(args_schema=QuestionAndAnswerInput)
def question_and_answer_tool(
    question: str,
    tool_call_id: Annotated[str, InjectedToolCallId],
    state: Annotated[dict, InjectedState],
) -> Dict[str, Any]:
    """
    Answer a question using PDF content stored in the state via retrieval augmented generation.

    This tool retrieves the PDF binary data from the state (under the key "pdf_data"), extracts its
    textual content, and generates an answer to the specified question. It also extracts the 
    llm_model (of type BaseChatModel) from the state to use for answering.

    Args:
        question (str): The question regarding the PDF content.
        tool_call_id (str): Unique identifier for the current tool call.
        state (dict): A dictionary representing the current state, expected to contain PDF data
                      under the key "pdf_data" with a sub-key "pdf_object" for the binary content,
                      and a key "llm_model" holding the language model instance.

    Returns:
        Dict[str, Any]: A dictionary containing the generated answer or an error message.
    """
    logger.info("Starting PDF Question and Answer tool using PDF data from state.")
    pdf_state = state.get("pdf_data")
    if not pdf_state:
        error_msg = "No pdf_data found in state."
        logger.error(error_msg)
        return Command(
            update={
                "messages": [
                    ToolMessage(content=error_msg, tool_call_id=tool_call_id)
                ]
            }
        )
    pdf_bytes = pdf_state.get("pdf_object")
    if not pdf_bytes:
        error_msg = "PDF binary data is missing in the pdf_data from state."
        logger.error(error_msg)
        return Command(
            update={
                "messages": [
                    ToolMessage(content=error_msg, tool_call_id=tool_call_id)
                ]
            }
        )
    # Retrieve llm_model from state; use a default if not provided.
    llm_model = state.get("llm_model")
    if not llm_model:
        logger.error("Missing LLM model instance in state.")
        return {"error": "No LLM model found in state."}
    answer = generate_answer(question, pdf_bytes, llm_model)
    logger.info("Generated answer: %s", answer)
    return answer
