# chat_logic.py

import os
import re
import warnings
from pathlib import Path
from typing import Any, Tuple, Optional, Dict

# Langchain/OpenAI imports
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
from langchain_classic.chains import ConversationalRetrievalChain
from langchain_classic.memory import ConversationBufferMemory, ConversationSummaryBufferMemory
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter, CharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.document_transformers import EmbeddingsRedundantFilter, LongContextReorder
from langchain_classic.retrievers.document_compressors import DocumentCompressorPipeline
from langchain_classic.retrievers.document_compressors import EmbeddingsFilter
from langchain_classic.retrievers import ContextualCompressionRetriever
from langchain_text_splitters import TextSplitter
from langchain_core.retrievers import BaseRetriever
from langchain_core.language_models import BaseChatModel

# --- Constants & Helpers ---

LOCAL_VECTOR_STORE_DIR = Path(__file__).resolve().parent.joinpath("data", "vector_stores")
# !!! SET YOUR DEFAULT PDF PATH HERE !!!
# Assuming the default PDF is in the same directory as this script.
DEFAULT_PDF_PATH = Path(__file__).resolve().parent.joinpath("S:\\ano_dec_pro\\AnomalyDetectionCVPR2018-Pytorch\\ring_chat_bot\\Ring_App_Documentation.pdf")
DEFAULT_VECTORSTORE_NAME = "default_pdf_db"
OPENAI_KEY = os.getenv("OPENAI_API_KEY")


def ensure_dir(p: Path) -> None:
    p.mkdir(parents=True, exist_ok=True)

def load_default_pdf():
    # Attempt to find the default PDF in the script directory
    if not DEFAULT_PDF_PATH.exists():
        raise FileNotFoundError(
            f"Default PDF not found: {DEFAULT_PDF_PATH}. Please place your PDF here or update the path in chat_logic.py"
        )
    loader = PyPDFLoader(DEFAULT_PDF_PATH.as_posix())
    return loader.load()

def split_documents(docs, chunk_size: int = 1600, chunk_overlap: int = 200):
    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    return splitter.split_documents(docs)

def select_embeddings(openai_key: str | None) -> OpenAIEmbeddings:
    if not openai_key:
        raise ValueError("OPENAI_API_KEY is required.")
    return OpenAIEmbeddings(api_key=openai_key)

# --- Core RAG Components ---

def vectorstore_backed_retriever(vs: Chroma, search_type: str = "similarity", k: int = 16, score_threshold: float | None = None) -> BaseRetriever:
    kwargs = {}
    if k is not None:
        kwargs["k"] = k
    if score_threshold is not None:
        kwargs["score_threshold"] = score_threshold
    return vs.as_retriever(search_type=search_type, search_kwargs=kwargs)


def make_compression_retriever(embeddings: OpenAIEmbeddings, base_retriever: BaseRetriever, chunk_size: int = 500, k: int = 16, similarity_threshold: float | None = None) -> ContextualCompressionRetriever:
    splitter: TextSplitter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=0, separator=". ")
    redundant_filter = EmbeddingsRedundantFilter(embeddings=embeddings)
    relevant_filter = EmbeddingsFilter(embeddings=embeddings, k=k, similarity_threshold=similarity_threshold)
    reordering = LongContextReorder()
    pipeline = DocumentCompressorPipeline(transformers=[splitter, redundant_filter, relevant_filter, reordering])
    return ContextualCompressionRetriever(base_compressor=pipeline, base_retriever=base_retriever)


def make_memory(model_name: str, openai_key: str | None):
    # Simplified memory logic for Streamlit
    return ConversationSummaryBufferMemory(
        max_token_limit=1024,
        llm=ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key=openai_key, temperature=0.1),
        return_messages=True,
        memory_key="chat_history",
        output_key="answer",
        input_key="question",
    )


def answer_template(language: str = "english") -> str:
    return f"""Answer the question at the end, using only the following context (delimited by <context></context>).
Your answer must be in the language at the end.

<context>
{{chat_history}}

{{context}}
</context>

Question: {{question}}

Language: {language}.
"""

def build_chain(model: str, retriever: BaseRetriever, openai_key: str | None) -> Tuple[ConversationalRetrievalChain, Any]:
    condense_question_prompt = PromptTemplate(
        input_variables=["chat_history", "question"],
        template=(
            "Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question, in its original language.\n\nChat History:\n{chat_history}\n\nFollow Up Input: {question}\n\nStandalone question:"
        ),
    )
    answer_prompt = ChatPromptTemplate.from_template(answer_template(language="english"))
    memory = make_memory(model, openai_key)

    standalone_llm = ChatOpenAI(api_key=openai_key, model=model, temperature=0.1)
    response_llm = ChatOpenAI(api_key=openai_key, model=model, temperature=0.5, top_p=0.95)

    chain = ConversationalRetrievalChain.from_llm(
        condense_question_prompt=condense_question_prompt,
        combine_docs_chain_kwargs={"prompt": answer_prompt},
        condense_question_llm=standalone_llm,
        llm=response_llm,
        memory=memory,
        retriever=retriever,
        chain_type="stuff",
        verbose=False,
        return_source_documents=True,
    )
    return chain, memory


def setup_default_rag(openai_key: str, model_name: str = "gpt-4-turbo") -> Tuple[ConversationalRetrievalChain, Any]:
    """
    Sets up the RAG chain using the default hardcoded PDF file.
    This replaces the file upload functionality for the initial setup.
    """
    
    vectorstore_path = LOCAL_VECTOR_STORE_DIR.joinpath(DEFAULT_VECTORSTORE_NAME)
    ensure_dir(vectorstore_path)
    
    embeddings = select_embeddings(openai_key)

    # Check if the vector store already exists locally (persistence logic)
    if not any(vectorstore_path.iterdir()):
        # 1. Load and split the default PDF
        docs = load_default_pdf()
        chunks = split_documents(docs)
        
        # 2. Create and persist the Vector Store (Chroma)
        vs = Chroma.from_documents(
            documents=chunks, 
            embedding=embeddings, 
            persist_directory=vectorstore_path.as_posix()
        )
        vs.persist()
    else:
        # 3. Load the existing Vector Store
        vs = Chroma(embedding_function=embeddings, persist_directory=vectorstore_path.as_posix())

    # 4. Create Retriever
    base_retriever = vectorstore_backed_retriever(vs)
    retriever = make_compression_retriever(embeddings=embeddings, base_retriever=base_retriever)

    # 5. Build and return chain
    chain, memory = build_chain(model_name, retriever, openai_key)

    return chain, memory

# The process_uploaded_file function is removed as we are hardcoding the default file setup.