app/rag.py

import os
import re
import unicodedata
from pathlib import Path
from typing import List

from dotenv import load_dotenv
import google.generativeai as genai
from huggingface_hub import InferenceClient

load_dotenv()

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
if GEMINI_API_KEY:
    genai.configure(api_key=GEMINI_API_KEY)

vectordb = None
retriever = None
embeddings = None
rag_initialized = False
uploaded_documents = []
last_index_mtime = None

RAG_DATA_DIR = Path(__file__).resolve().parent.parent / "rag_data"
FAISS_INDEX_PATH = RAG_DATA_DIR / "faiss_index"
INSUFFICIENT_CONTEXT_MARKER = "i don't have enough information in the documents"


def initialize_embeddings():
    """Initialize the multilingual embedding model."""
    global embeddings

    if embeddings is not None:
        return embeddings

    print("Loading multilingual embedding model...")
    from langchain_huggingface import HuggingFaceEmbeddings

    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
        encode_kwargs={"normalize_embeddings": True},
    )
    print("Embedding model loaded.")
    return embeddings


def clean_text(text: str) -> str:
    """Clean and normalize text for embedding."""
    if not isinstance(text, str) or not text.strip():
        return ""

    normalized_text = unicodedata.normalize("NFKC", text)

    cleaned_chars = [
        char for char in normalized_text
        if unicodedata.category(char) not in ["So", "Cn", "Cc", "Cf", "Cs"]
    ]
    cleaned_text = "".join(cleaned_chars)

    cleaned_text = re.sub(r"\s+", " ", cleaned_text).strip()

    return cleaned_text


def load_and_process_pdf(pdf_path: str) -> List[dict]:
    """Load a PDF and split it into chunks."""
    from langchain_community.document_loaders import PyPDFLoader
    from langchain_text_splitters import RecursiveCharacterTextSplitter

    print(f"Loading PDF: {pdf_path}")

    loader = PyPDFLoader(pdf_path)
    docs = loader.load()

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=300,
        chunk_overlap=80,
    )
    chunks = splitter.split_documents(docs)

    print(f"Loaded {len(docs)} pages, created {len(chunks)} chunks.")
    return chunks


def create_vector_store(chunks: List) -> bool:
    """Create or update the FAISS vector store with document chunks."""
    global vectordb, retriever, rag_initialized

    from langchain_community.vectorstores import FAISS

    initialize_embeddings()

    texts = [doc.page_content for doc in chunks]
    metadatas = [doc.metadata for doc in chunks]

    processed_texts = []
    processed_metadatas = []

    for i, text in enumerate(texts):
        cleaned_text = clean_text(text)
        if cleaned_text:
            processed_texts.append(cleaned_text)
            processed_metadatas.append(metadatas[i])

    if not processed_texts:
        print("No valid texts after cleaning.")
        return False

    print(f"Processing {len(processed_texts)} text chunks for embedding...")

    if vectordb is None:
        vectordb = FAISS.from_texts(processed_texts, embeddings, metadatas=processed_metadatas)
    else:
        new_vectordb = FAISS.from_texts(processed_texts, embeddings, metadatas=processed_metadatas)
        vectordb.merge_from(new_vectordb)

    retriever = vectordb.as_retriever(search_kwargs={"k": 4})
    rag_initialized = True

    save_vector_store()
    _sync_uploaded_documents()

    print("Vector store created/updated successfully.")
    return True


def save_vector_store():
    """Save the FAISS index to disk."""
    global vectordb, last_index_mtime

    if vectordb is None:
        return

    RAG_DATA_DIR.mkdir(parents=True, exist_ok=True)
    vectordb.save_local(str(FAISS_INDEX_PATH))
    last_index_mtime = _get_index_mtime()
    print(f"Vector store saved to {FAISS_INDEX_PATH}.")


def load_vector_store() -> bool:
    """Load the FAISS index from disk if it exists."""
    global vectordb, retriever, rag_initialized, last_index_mtime

    if not FAISS_INDEX_PATH.exists():
        return False

    try:
        from langchain_community.vectorstores import FAISS

        initialize_embeddings()
        vectordb = FAISS.load_local(
            str(FAISS_INDEX_PATH),
            embeddings,
            allow_dangerous_deserialization=True,
        )
        retriever = vectordb.as_retriever(search_kwargs={"k": 4})
        rag_initialized = True
        last_index_mtime = _get_index_mtime()
        _sync_uploaded_documents()
        print("Loaded existing vector store from disk.")
        return True
    except Exception as e:
        print(f"Failed to load vector store: {e}")
        return False


def rag_answer(question: str) -> dict:
    """Answer a question using RAG - first check database, then fallback to Gemini/HF."""
    global retriever, vectordb, last_index_mtime

    result = {
        "answer": "",
        "source": "none",
        "context_found": False,
        "relevance_score": 0.0,
    }

    if FAISS_INDEX_PATH.exists():
        current_mtime = _get_index_mtime()
        if (not rag_initialized or retriever is None) or (
            current_mtime and last_index_mtime and current_mtime > last_index_mtime
        ):
            load_vector_store()

    if not rag_initialized or retriever is None:
        result["source"] = "gemini"
        result["answer"] = _ask_gemini_directly(question)
        return result

    docs_with_scores = vectordb.similarity_search_with_score(question, k=4)

    if not docs_with_scores:
        print(f"No documents found for question: {question}")
        result["source"] = "gemini"
        result["answer"] = _ask_gemini_directly(question)
        return result

    best_score = docs_with_scores[0][1] if docs_with_scores else float("inf")
    result["relevance_score"] = float(best_score)

    print(f"\nQuestion: {question}")
    print(f"Retrieved {len(docs_with_scores)} documents:")
    for i, (doc, score) in enumerate(docs_with_scores):
        preview = doc.page_content[:100].replace("\n", " ")
        print(f"  [{i + 1}] Score: {score:.3f} - {preview}...")

    print(f"Using RAG with relevance score: {best_score}")

    docs = [doc for doc, score in docs_with_scores]
    context = "\n\n".join([d.page_content for d in docs])
    result["context_found"] = True

    prompt = (
        "You are a helpful assistant. Answer the question based ONLY on the following "
        "context from the PDF document. If the context doesn't contain enough information "
        "to answer the question, say \"I don't have enough information in the documents to "
        "answer this question.\"\n\n"
        "Context from PDF:\n"
        f"{context}\n\n"
        f"Question: {question}\n\n"
        "Answer (in English):"
    )

    try:
        gemini_key = os.getenv("GEMINI_API_KEY")
        if gemini_key:
            try:
                model = genai.GenerativeModel("models/gemini-2.5-flash")
                response = model.generate_content(prompt)
                rag_answer_text = (response.text or "").strip()
                if _is_insufficient_context_answer(rag_answer_text):
                    print("RAG context not sufficient. Falling back to direct AI answer.")
                    result["answer"] = _ask_gemini_directly(question)
                    result["source"] = "gemini"
                    return result
                result["answer"] = rag_answer_text
                result["source"] = "rag"
                return result
            except Exception as gemini_error:
                error_msg = str(gemini_error)
                print(f"Gemini error in RAG: {error_msg[:200]}...")
                if "429" in error_msg or "quota" in error_msg.lower():
                    print("Gemini quota exceeded. Using Hugging Face for RAG.")

        print("Using Hugging Face for RAG answer...")
        rag_answer_text = _ask_huggingface_free(prompt).strip()
        if _is_insufficient_context_answer(rag_answer_text):
            print("RAG context not sufficient. Falling back to direct AI answer.")
            result["answer"] = _ask_gemini_directly(question)
            result["source"] = "gemini"
            return result
        result["answer"] = rag_answer_text
        result["source"] = "rag"

    except Exception as e:
        print(f"All RAG generation failed: {e}")
        result["answer"] = "Sorry, unable to generate answer. Please try again later."
        result["source"] = "error"

    return result


def _ask_huggingface_free(prompt: str) -> str:
    """Use free Hugging Face Inference API with token if available."""
    hf_token = os.getenv("HF_API_TOKEN")

    try:
        client = InferenceClient(token=hf_token)
    except Exception as e:
        raise Exception(f"Failed to create Hugging Face client: {e}")

    messages = [{"role": "user", "content": prompt}]

    try:
        print("Calling Hugging Face API (Qwen2.5-72B-Instruct)...")
        response = client.chat_completion(
            messages=messages,
            model="Qwen/Qwen2.5-72B-Instruct",
            max_tokens=500,
            temperature=0.7,
        )
        return response.choices[0].message.content
    except Exception as e:
        error_str = str(e)
        print(f"Hugging Face primary model error: {e}")

        try:
            print("Trying backup model (Microsoft Phi-3)...")
            response = client.chat_completion(
                messages=messages,
                model="microsoft/Phi-3-mini-4k-instruct",
                max_tokens=500,
                temperature=0.7,
            )
            return response.choices[0].message.content
        except Exception as e2:
            print(f"Backup model also failed: {e2}")
            raise Exception(f"All HF models failed: {error_str}")


def _ask_gemini_directly(question: str) -> str:
    """Fallback: Ask Gemini directly without RAG context, with Hugging Face fallback."""
    prompt = (
        "Answer the following question helpfully and accurately:\n\n"
        f"Question: {question}\n\n"
        "Answer:"
    )

    gemini_key = os.getenv("GEMINI_API_KEY")

    if gemini_key:
        try:
            model = genai.GenerativeModel("models/gemini-2.5-flash")
            response = model.generate_content(prompt)
            return response.text
        except Exception as gemini_error:
            error_msg = str(gemini_error)
            print(f"Gemini API error: {error_msg[:200]}...")

            if "429" in error_msg or "quota" in error_msg.lower():
                print("Gemini quota exceeded. Switching to Hugging Face.")
            else:
                print("Gemini error. Switching to Hugging Face.")
    else:
        print("No Gemini API key, using Hugging Face.")

    try:
        print("Using Hugging Face for direct answer...")
        return _ask_huggingface_free(prompt)
    except Exception as hf_error:
        print(f"Hugging Face error: {hf_error}")
        return (
            "Sorry, both AI services are unavailable. "
            f"Gemini quota exceeded, and Hugging Face error: {str(hf_error)}"
        )


def get_rag_status() -> dict:
    """Get the current status of the RAG system."""
    if not rag_initialized and FAISS_INDEX_PATH.exists():
        load_vector_store()

    _sync_uploaded_documents()
    return {
        "initialized": rag_initialized,
        "documents_count": len(uploaded_documents),
        "documents": uploaded_documents,
        "has_embeddings": embeddings is not None,
        "has_vector_store": vectordb is not None,
    }


def clear_rag_data():
    """Clear all RAG data."""
    global vectordb, retriever, rag_initialized, uploaded_documents, last_index_mtime

    vectordb = None
    retriever = None
    rag_initialized = False
    uploaded_documents = []
    last_index_mtime = None

    if FAISS_INDEX_PATH.exists():
        import shutil

        shutil.rmtree(FAISS_INDEX_PATH)

    print("RAG data cleared.")
    return True


def _get_index_mtime():
    index_file = FAISS_INDEX_PATH / "index.faiss"
    if index_file.exists():
        return index_file.stat().st_mtime
    return None


def _is_insufficient_context_answer(answer_text: str) -> bool:
    if not answer_text:
        return True
    normalized = answer_text.strip().lower()
    return INSUFFICIENT_CONTEXT_MARKER in normalized


def _sync_uploaded_documents():
    global uploaded_documents

    if not RAG_DATA_DIR.exists():
        uploaded_documents = []
        return

    uploaded_documents = sorted(
        [pdf.name for pdf in RAG_DATA_DIR.glob("*.pdf") if pdf.is_file()]
    )


def rebuild_vector_store_from_pdfs() -> bool:
    """Rebuild vector store from all PDFs in rag_data directory."""
    global vectordb, retriever, rag_initialized

    _sync_uploaded_documents()
    if not uploaded_documents:
        print("No PDFs found in rag_data to rebuild vector store.")
        return False

    initialize_embeddings()

    vectordb = None
    retriever = None
    rag_initialized = False

    all_chunks = []
    for filename in uploaded_documents:
        pdf_path = RAG_DATA_DIR / filename
        try:
            chunks = load_and_process_pdf(str(pdf_path))
            all_chunks.extend(chunks)
        except Exception as e:
            print(f"Skipping PDF '{filename}' due to processing error: {e}")

    if not all_chunks:
        print("No chunks generated from PDFs. Rebuild aborted.")
        return False

    success = create_vector_store(all_chunks)
    if success:
        print(f"Rebuilt vector store from {len(uploaded_documents)} PDF(s).")
    return success