import os
import pickle
import json
import logging
import re
import time
import urllib.error
import urllib.request
from functools import lru_cache
from typing import Any, Iterator, List
from dotenv import load_dotenv
from pinecone import Pinecone

from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain_community.retrievers import BM25Retriever
from langchain.chains import RetrievalQA
from langchain_core.prompts import PromptTemplate
from langchain_core.retrievers import BaseRetriever
from langchain_core.callbacks import CallbackManagerForRetrieverRun

load_dotenv()

# ===============================
# CONFIG
# ===============================
INDEX_NAME = "branham-index"
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CHUNKS_FILE = os.path.join(BASE_DIR, "sermon_chunks.pkl")
GEMINI_EMBEDDING_MODEL = "models/gemini-embedding-001"
PINECONE_DIMENSION = 768
PRIORITY_K = 8
SERIES_DOCS_PER_SERMON = 2
BM25_K = 8
VECTOR_K = 8
MAX_CONTEXT_DOCS = 20
MIN_ENTITY_LENGTH = 4
LLM_PROVIDER = os.getenv("LLM_PROVIDER", "groq").strip().lower()
GROQ_MODEL = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
GEMINI_MODEL = os.getenv("GEMINI_MODEL", "gemini-2.5-flash")
logger = logging.getLogger(__name__)
UNSUPPORTED_ANSWER = "The retrieved sermon excerpts do not give enough support to answer that safely."

CONTEXT_POLICY = {
    "Strict Mode": {
        "max_docs": 8,
        "selection": "ranked",
    },
    "Study Mode": {
        "max_docs": 14,
        "selection": "ranked",
    },
    "Summary Mode": {
        "max_docs": 20,
        "selection": "distributed",
    },
}


# ===============================
# CANONICAL SERIES
# ===============================
SEVEN_SEALS_CANON = [
    "63-0317M God Hiding Himself In Simplicity, Then Revealing Himself In The Same.pdf",
    "63-0317E The Breach Between The Seven Church Ages And The Seven Seals.pdf",
    "63-0318 The First Seal.pdf",
    "63-0319 The Second Seal.pdf",
    "63-0320 The Third Seal.pdf",
    "63-0321 The Fourth Seal.pdf",
    "63-0322 The Fifth Seal.pdf",
    "63-0323 The Sixth Seal.pdf",
    "63-0324M Questions And Answers On The Seals.pdf",
    "63-0324E The Seventh Seal.pdf",
]

ADOPTION_CANON = [
    "60-0515E Adoption #1.pdf",
    "60-0518 Adoption #2.pdf",
    "60-0522M Adoption #3.pdf",
    "60-0522E Adoption #4.pdf",
]

DEMONOLOGY_CANON = [
    "53-0608A Demonology, Physical Realm.pdf",
    "53-0609A Demonology, Religious Realm.pdf",
    "55-0724 Enticing Spirits.pdf",
]

SERIES_GROUPS = {
    "seven seals": SEVEN_SEALS_CANON,
    "adoption": ADOPTION_CANON,
    "demonology": DEMONOLOGY_CANON,
}

SEAL_SERMONS = [
    "63-0318 The First Seal.pdf",
    "63-0319 The Second Seal.pdf",
    "63-0320 The Third Seal.pdf",
    "63-0321 The Fourth Seal.pdf",
    "63-0322 The Fifth Seal.pdf",
    "63-0323 The Sixth Seal.pdf",
    "63-0324E The Seventh Seal.pdf",
]

SINGLE_SEAL_ALIASES = {
    "first seal": "63-0318 The First Seal.pdf",
    "1st seal": "63-0318 The First Seal.pdf",
    "second seal": "63-0319 The Second Seal.pdf",
    "2nd seal": "63-0319 The Second Seal.pdf",
    "third seal": "63-0320 The Third Seal.pdf",
    "3rd seal": "63-0320 The Third Seal.pdf",
    "fourth seal": "63-0321 The Fourth Seal.pdf",
    "4th seal": "63-0321 The Fourth Seal.pdf",
    "fifth seal": "63-0322 The Fifth Seal.pdf",
    "5th seal": "63-0322 The Fifth Seal.pdf",
    "sixth seal": "63-0323 The Sixth Seal.pdf",
    "6th seal": "63-0323 The Sixth Seal.pdf",
    "seventh seal": "63-0324E The Seventh Seal.pdf",
    "7th seal": "63-0324E The Seventh Seal.pdf",
}

SEVEN_SEALS_QUERY_HINTS = {
    "63-0318": "first seal white horse bow crown conquer",
    "63-0319": "second seal red horse sword take peace kill",
    "63-0320": "third seal black horse balances wheat barley oil wine",
    "63-0321": "fourth seal pale horse death hell eagle",
    "63-0322": "fifth seal souls under the altar white robes Jews",
    "63-0323": "sixth seal earthquake sun black moon blood stars fall",
    "63-0324E": "seventh seal silence seven thunders coming of Christ end time",
}

SERIES_QUERY_HINTS = {
    **SEVEN_SEALS_QUERY_HINTS,
    "60-0515E": "adoption placing sons predestination inheritance position",
    "60-0518": "adoption sons inheritance child placing position",
    "60-0522M": "adoption holy spirit church position sons",
    "60-0522E": "adoption manifestation sons inheritance placing",
    "53-0608A": "demonology physical realm spirits sickness oppression demons",
    "53-0609A": "demonology religious realm spirits worship deception demons",
    "55-0724": "enticing spirits deception religious spirits influence",
}

ANSWER_MODE_INSTRUCTIONS = {
    "Strict Mode": (
        "Use strict grounding. If the retrieved excerpts do not directly support the answer, "
        f"respond exactly: {UNSUPPORTED_ANSWER}"
    ),
    "Study Mode": (
        "Give a fuller study-style explanation, but keep every doctrinal claim tied to the "
        "provided excerpts and source references."
    ),
    "Summary Mode": (
        "Prioritize concise summaries. If the user asks about a sermon, summarize only what the "
        "retrieved excerpts from that sermon support."
    ),
}

# ===============================
# HELPERS
# ===============================
class GeminiEmbedding768(Embeddings):
    """Custom local HuggingFace BGE embedding client matched to the 768-dimension index."""

    def __init__(self) -> None:
        logger.info("Initializing local HuggingFace BGE embedding client (BAAI/bge-base-en-v1.5)...")
        # Load HuggingFace BGE model locally on CPU (or CUDA/MPS if available, default to CPU for general reliability)
        self.embeddings = HuggingFaceEmbeddings(
            model_name="BAAI/bge-base-en-v1.5",
            model_kwargs={'device': 'cpu'},
            encode_kwargs={'normalize_embeddings': True}
        )
        logger.info("Local BGE embedding client initialized successfully!")

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        if not texts:
            return []
        return self.embeddings.embed_documents(texts)

    def embed_query(self, text: str) -> List[float]:
        return self.embeddings.embed_query(text)


def normalize(text: str) -> str:
    return text.lower().replace("_", " ").replace("-", " ").strip()

@lru_cache(maxsize=1)
def load_chunks() -> List[Document]:
    if not os.path.exists(CHUNKS_FILE):
        return []
    with open(CHUNKS_FILE, "rb") as f:
        return pickle.load(f)


@lru_cache(maxsize=1)
def get_bm25_retriever():
    bm25 = BM25Retriever.from_documents(load_chunks())
    bm25.k = BM25_K
    return bm25


@lru_cache(maxsize=1)
def get_search_bm25_retriever():
    bm25 = BM25Retriever.from_documents(load_chunks())
    bm25.k = 50
    return bm25


@lru_cache(maxsize=1)
def get_vector_store():
    return PineconeVectorStore(
        index_name=INDEX_NAME,
        embedding=GeminiEmbedding768(),
    )


def extract_date_code(filename: str) -> str:
    """
    Assumes filenames start with NN-NNNNE
    Example: 62-0909E In His Presence.pdf
    """
    return filename.split()[0].replace(".pdf", "")


def extract_query_sermon_code(query: str) -> str | None:
    match = re.search(r"\b\d{2}-\d{4}[A-Z]?\b", query.upper())
    return match.group(0) if match else None


def source_matches_code(source: str, code: str) -> bool:
    return extract_date_code(source).upper() == code.upper()


def source_matches_title(source: str, title: str) -> bool:
    return source_matches_code(source, extract_date_code(title).upper())


def source_matches_any_code(source: str, codes: set[str]) -> bool:
    return extract_date_code(source).upper() in codes


def document_key(doc: Document) -> tuple[str, str, str]:
    return (
        doc.metadata.get("source", ""),
        str(doc.metadata.get("paragraph", "")),
        doc.page_content[:120],
    )


def messagehub_link(filename: str) -> str:
    code = extract_date_code(filename)
    return f"https://www.messagehub.info/en/read.do?ref_num={code}"


STOPWORDS = {
    "the", "a", "an", "of", "in", "on", "at", "and", "to", "for", "with", "by",
    "what", "who", "where", "when", "why", "how", "did", "does", "do", "is",
    "are", "was", "were", "be", "been", "being", "about", "teach", "taught",
    "said", "say", "says", "explain", "explained", "brother", "branham",
    "william", "sermon", "sermons", "message", "messages", "summarize",
    "summary", "list", "state", "meaning", "events", "relationship",
}
GENERIC_TITLE_TOKENS = {
    "know", "believe", "faith", "love", "life", "way", "truth", "word",
    "god", "lord", "jesus", "christ", "christian", "church", "message",
}


def normalize_text(text: str) -> str:
    text = text.lower()
    text = re.sub(r"[^a-z0-9\s]", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def extract_sermon_title(filename: str) -> str:
    """
    '62-0909E In His Presence.pdf' → 'in his presence'
    """
    name = filename.replace(".pdf", "").replace(".PDF", "")

    parts = name.split(" ", 1)
    if len(parts) == 2 and "-" in parts[0]:
        name = parts[1]

    return normalize_text(name)


def tokenize_meaningful(text: str) -> set:
    return {
        w for w in normalize_text(text).split()
        if w not in STOPWORDS and len(w) > 2
    }


def sermon_title_matches(user_query: str, filename: str) -> bool:
    """
    Match only if ALL meaningful title words exist in user query.
    Prevents partial matches like 'presence'.
    """
    title_tokens = tokenize_meaningful(extract_sermon_title(filename))
    query_tokens = tokenize_meaningful(user_query)

    if not title_tokens:
        return False

    return title_tokens.issubset(query_tokens)


def explicit_title_lookup(query: str) -> bool:
    query_clean = normalize_text(query)
    if extract_query_sermon_code(query):
        return True
    title_lookup_phrases = {
        "sermon called",
        "message called",
        "sermon titled",
        "message titled",
        "sermon named",
        "message named",
        "the sermon",
        "the message",
        "summarize",
        "summary of",
    }
    return any(phrase in query_clean for phrase in title_lookup_phrases)


def safe_sermon_title_matches(user_query: str, filename: str) -> bool:
    title_tokens = tokenize_meaningful(extract_sermon_title(filename))
    if not title_tokens:
        return False
    if len(title_tokens) == 1 and next(iter(title_tokens)) in GENERIC_TITLE_TOKENS:
        return False
    if len(title_tokens) < 2 and not explicit_title_lookup(user_query):
        return False
    return sermon_title_matches(user_query, filename)


def detect_single_seal_title(query: str) -> str | None:
    query_clean = normalize_text(query)
    for alias, title in SINGLE_SEAL_ALIASES.items():
        if alias in query_clean:
            return title
    return None


def query_entity_tokens(query: str) -> set[str]:
    return {
        token for token in tokenize_meaningful(query)
        if len(token) >= MIN_ENTITY_LENGTH
    }


def rank_by_query_terms(docs: List[Document], query: str) -> List[Document]:
    terms = query_entity_tokens(query)

    def score(doc: Document) -> tuple[int, int]:
        text = normalize_text(doc.page_content)
        hits = sum(1 for term in terms if term in text)
        return hits, -len(text)

    return sorted(docs, key=score, reverse=True)


def distributed_sample(docs: List[Document], limit: int) -> List[Document]:
    if limit <= 0 or len(docs) <= limit:
        return docs[:limit]
    if limit == 1:
        return [docs[0]]

    step = (len(docs) - 1) / (limit - 1)
    selected = []
    used_indexes = set()
    for i in range(limit):
        index = round(i * step)
        if index not in used_indexes:
            selected.append(docs[index])
            used_indexes.add(index)
    return selected


def select_context_docs(docs: List[Document], query: str, answer_mode: str) -> List[Document]:
    policy = CONTEXT_POLICY.get(answer_mode, CONTEXT_POLICY["Strict Mode"])
    max_docs = policy["max_docs"]

    if "seven seals" in normalize(query):
        guaranteed_titles = SEVEN_SEALS_CANON if answer_mode == "Summary Mode" else SEAL_SERMONS
        selected = []
        seen = set()
        for title in guaranteed_titles:
            code = extract_date_code(title).upper()
            for doc in docs:
                if source_matches_code(doc.metadata.get("source", ""), code):
                    key = document_key(doc)
                    if key not in seen:
                        selected.append(doc)
                        seen.add(key)
                    break
        remaining = max_docs - len(selected)
        if remaining > 0:
            for doc in docs:
                key = document_key(doc)
                if key not in seen:
                    selected.append(doc)
                    seen.add(key)
                    remaining -= 1
                if remaining <= 0:
                    break
        return selected[:max_docs]

    if policy["selection"] == "distributed":
        priority_docs = [d for d in docs if d.metadata.get("retrieval_origin") == "priority"]
        other_docs = [d for d in docs if d.metadata.get("retrieval_origin") != "priority"]
        selected = distributed_sample(priority_docs, min(len(priority_docs), max_docs))
        remaining = max_docs - len(selected)
        if remaining > 0:
            selected.extend(rank_by_query_terms(other_docs, query)[:remaining])
        return selected[:max_docs]

    priority_docs = [d for d in docs if d.metadata.get("retrieval_origin") == "priority"]
    other_docs = [d for d in docs if d.metadata.get("retrieval_origin") != "priority"]
    selected = priority_docs[: min(len(priority_docs), max_docs)]
    remaining = max_docs - len(selected)
    if remaining > 0:
        selected.extend(rank_by_query_terms(other_docs, query)[:remaining])
    return selected[:max_docs]


def query_support_terms(query: str) -> set[str]:
    terms = query_entity_tokens(query)
    return {
        term for term in terms
        if term not in STOPWORDS and not re.fullmatch(r"\d+", term)
    }


def evaluate_retrieval_confidence(
    query: str,
    docs: List[Document],
    diagnostics: dict[str, Any],
) -> dict[str, Any]:
    terms = query_support_terms(query)
    context = normalize_text(" ".join(doc.page_content for doc in docs))
    supported_terms = sorted(term for term in terms if term in context)
    missing_terms = sorted(terms - set(supported_terms))
    coverage = len(supported_terms) / len(terms) if terms else 1.0

    confidence = "supported"
    reason = "priority_or_supported_terms"
    should_refuse = False

    if not docs:
        confidence = "insufficient"
        reason = "no_documents"
        should_refuse = True
    elif diagnostics.get("priority_count", 0) > 0:
        confidence = "supported"
        reason = "priority_match"
    elif terms and missing_terms:
        confidence = "insufficient"
        reason = "missing_query_terms_without_priority_match"
        should_refuse = True
    elif diagnostics.get("source_count", 0) < 2:
        confidence = "partial"
        reason = "too_few_documents"

    return {
        "retrieval_confidence": confidence,
        "confidence_reason": reason,
        "should_refuse_generation": should_refuse,
        "query_terms": sorted(terms),
        "supported_query_terms": supported_terms,
        "missing_query_terms": missing_terms,
        "query_term_coverage": round(coverage, 3),
    }

# ===============================
# RETRIEVER
# ===============================
class BranhamRetriever(BaseRetriever):
    """
    NotebookLM-style hybrid retriever:
    - local priority
    - semantic fallback
    - series-aware
    - safe + deduplicated
    """

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun = None
    ) -> List[Document]:
        docs, _ = retrieve_documents_with_diagnostics(query, answer_mode="Strict Mode")
        return docs


def retrieve_documents_with_diagnostics(query: str, answer_mode: str = "Strict Mode"):
    started_at = time.perf_counter()
    query_clean = normalize(query)
    chunks = load_chunks()
    results: List[Document] = []
    seen = set()
    counts = {
        "priority": 0,
        "bm25": 0,
        "pinecone": 0,
    }
    pinecone_status = "not_run"
    pinecone_error = ""

    def add_doc(doc: Document, origin: str) -> bool:
        key = document_key(doc)
        if key in seen:
            return False
        metadata = dict(doc.metadata)
        metadata["retrieval_origin"] = origin
        metadata["retrieval_rank"] = len(results) + 1
        results.append(Document(page_content=doc.page_content, metadata=metadata))
        seen.add(key)
        counts[origin] = counts.get(origin, 0) + 1
        return True

    explicit_sermon = extract_query_sermon_code(query)
    single_seal_title = detect_single_seal_title(query)
    target_titles = []

    for key, titles in SERIES_GROUPS.items():
        if key in query_clean:
            target_titles = titles
            break

    if explicit_sermon:
        sermon_chunks = [
            d for d in chunks
            if source_matches_code(d.metadata.get("source", ""), explicit_sermon)
        ]
        if sermon_chunks:
            if answer_mode == "Summary Mode":
                for d in distributed_sample(sermon_chunks, CONTEXT_POLICY["Summary Mode"]["max_docs"]):
                    add_doc(d, "priority")
            else:
                sermon_bm25 = BM25Retriever.from_documents(sermon_chunks)
                sermon_bm25.k = PRIORITY_K
                for d in sermon_bm25.invoke(query):
                    add_doc(d, "priority")

    elif single_seal_title:
        target_code = extract_date_code(single_seal_title).upper()
        sermon_chunks = [
            d for d in chunks
            if source_matches_title(d.metadata.get("source", ""), single_seal_title)
        ]
        if sermon_chunks:
            seal_query = SERIES_QUERY_HINTS.get(target_code, query)
            for d in rank_by_query_terms(sermon_chunks, seal_query)[:PRIORITY_K]:
                add_doc(d, "priority")

    elif target_titles:
        if "seven seals" in query_clean and answer_mode != "Summary Mode":
            titles_for_retrieval = SEAL_SERMONS
        else:
            titles_for_retrieval = target_titles
        for title in titles_for_retrieval:
            target_code = extract_date_code(title).upper()
            sermon_chunks = [
                d for d in chunks
                if source_matches_code(d.metadata.get("source", ""), target_code)
            ]
            if sermon_chunks:
                series_query = SERIES_QUERY_HINTS.get(target_code, query)
                for d in rank_by_query_terms(sermon_chunks, series_query)[:SERIES_DOCS_PER_SERMON]:
                    add_doc(d, "priority")

    else:
        title_chunks = [
            d for d in chunks
            if safe_sermon_title_matches(query, d.metadata.get("source", ""))
        ]
        if title_chunks:
            if answer_mode == "Summary Mode":
                for d in distributed_sample(title_chunks, CONTEXT_POLICY["Summary Mode"]["max_docs"]):
                    add_doc(d, "priority")
            else:
                for d in rank_by_query_terms(title_chunks, query)[:PRIORITY_K]:
                    add_doc(d, "priority")

    max_context_docs = CONTEXT_POLICY.get(answer_mode, CONTEXT_POLICY["Strict Mode"])["max_docs"]
    if counts.get("priority", 0) < max_context_docs:
        for d in get_bm25_retriever().invoke(query):
            add_doc(d, "bm25")

    pre_vector_docs = select_context_docs(results, query, answer_mode)
    pre_vector_diagnostics = {
        "source_count": len(pre_vector_docs),
        "priority_count": sum(
            1 for d in pre_vector_docs if d.metadata.get("retrieval_origin") == "priority"
        ),
    }
    pre_vector_confidence = evaluate_retrieval_confidence(
        query,
        pre_vector_docs,
        pre_vector_diagnostics,
    )

    if counts.get("priority", 0) > 0:
        pinecone_status = "skipped_priority_match"
    elif pre_vector_confidence.get("should_refuse_generation"):
        pinecone_status = "skipped_low_confidence"
    else:
        try:
            vec_docs = get_vector_store().as_retriever(
                search_kwargs={"k": VECTOR_K}
            ).invoke(query)
            pinecone_status = "ok"
            for d in vec_docs:
                add_doc(d, "pinecone")
        except Exception as exc:
            pinecone_status = "failed"
            pinecone_error = str(exc)
            logger.warning("Pinecone vector search failed: %s", exc)

    final_docs = select_context_docs(results, query, answer_mode)
    distinct_sermons = {
        d.metadata.get("source", "")
        for d in final_docs
        if d.metadata.get("source")
    }
    diagnostics = {
        "retrieval_latency": round(time.perf_counter() - started_at, 3),
        "source_count": len(final_docs),
        "distinct_sermon_count": len(distinct_sermons),
        "sermons": sorted(distinct_sermons),
        "priority_count": sum(1 for d in final_docs if d.metadata.get("retrieval_origin") == "priority"),
        "bm25_count": sum(1 for d in final_docs if d.metadata.get("retrieval_origin") == "bm25"),
        "pinecone_count": sum(1 for d in final_docs if d.metadata.get("retrieval_origin") == "pinecone"),
        "pinecone_status": pinecone_status,
        "pinecone_error": pinecone_error,
        "context_policy": answer_mode,
        "candidate_count": len(results),
        "max_context_docs": max_context_docs,
    }
    diagnostics.update(evaluate_retrieval_confidence(query, final_docs, diagnostics))

    return final_docs, diagnostics


class _LegacyBranhamRetriever(BaseRetriever):
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun = None
    ) -> List[Document]:

        docs, _ = retrieve_documents_with_diagnostics(query, answer_mode="Strict Mode")
        return docs


# ===============================
# PROMPT
# ===============================
PROMPT_TEMPLATE = """
You are a careful research assistant helping one person understand William Marrion Branham's sermons.
Do not roleplay, impersonate, or speak as William Marrion Branham.
Always refer to him in the third person as "Brother Branham" or "William Branham."

RULES:
- You are speaking to only one person
- Be faithful to the sermons provided.
- Do NOT invent doctrine.
- Answer only from the provided context. Do not use outside knowledge.
- If the context does not clearly support the answer, say: "The retrieved sermon excerpts do not give enough support to answer that safely."
- If the context is partial, say what is supported and what is not supported.
- Do not smooth over contradictions or gaps. State them plainly.
- Use a respectful explanatory tone, not a preaching or prophetic persona.
- Be structured and clear.
- Use headings and bullet points.
- Explain symbols plainly.
- Prefer paraphrase, but preserve meaning.
- Ignore tape noise or filler language.
- If a question asks for a sermon summary, summarize only that sermon.
- If the question references the Seven Seals, prioritize the 1963 series.
- Phrase claims as "Brother Branham said/taught/explained..." when the context supports them.
- Do not say "I testified," "my ministry," "my dear friend," or anything that makes the assistant sound like Brother Branham.
- If context uses first-person sermon language, convert it to third-person attribution.
- Include source references for important claims using this format: [source, Para paragraph].
- Do not cite a source unless it appears in the provided context.
- Keep direct quotes short. Prefer paraphrase with source references.

ANSWER MODE:
{answer_mode_instructions}

CONTEXT:
{context_str}

QUESTION:
{question}

ANSWER:
"""

CONVERSATIONAL_PROMPT_TEMPLATE = """
You are a careful research assistant helping one person understand William Marrion Branham's sermons.
Do not roleplay, impersonate, or speak as William Marrion Branham.
Always refer to him in the third person as "Brother Branham" or "William Branham."

RULES:
- Be faithful to the sermons provided.
- Use the conversation history only to understand what the user's latest question refers to.
- Answer only from the provided context. Do not use outside knowledge.
- If the context does not clearly support the answer, say: "The retrieved sermon excerpts do not give enough support to answer that safely."
- If the context is partial, say what is supported and what is not supported.
- Use a natural conversational tone, but do not preach or impersonate.
- Phrase claims as "Brother Branham said/taught/explained..." when the context supports them.
- Do not say "I testified," "my ministry," "my dear friend," or anything that makes the assistant sound like Brother Branham.
- Include source references for important claims using this format: [source, Para paragraph].
- Do not cite a source unless it appears in the provided context.
- Keep direct quotes short. Prefer paraphrase with source references.

ANSWER MODE:
{answer_mode_instructions}

CONVERSATION HISTORY:
{conversation_history}

CONTEXT:
{context_str}

LATEST QUESTION:
{question}

ANSWER:
"""

QUESTION_REWRITE_PROMPT_TEMPLATE = """
Rewrite the user's latest question as a standalone search query for retrieving William Branham sermon excerpts.
Use the conversation history only to resolve references like "it", "that", "he", "those", "the first one", or "explain more".
Do not answer the question.
Do not add doctrine or facts not present in the conversation.
Return only the standalone question.

CONVERSATION HISTORY:
{conversation_history}

LATEST QUESTION:
{question}

STANDALONE QUESTION:
"""

PROMPT = PromptTemplate(
    template=PROMPT_TEMPLATE,
    input_variables=["context_str", "question"],
    partial_variables={
        "answer_mode_instructions": ANSWER_MODE_INSTRUCTIONS["Strict Mode"],
    },
)

DOCUMENT_PROMPT = PromptTemplate(
    template=(
        "SOURCE: {source}\n"
        "PARAGRAPH: {paragraph}\n"
        "TEXT:\n{page_content}"
    ),
    input_variables=["page_content", "source", "paragraph"],
)

QUESTION_REWRITE_PROMPT = PromptTemplate(
    template=QUESTION_REWRITE_PROMPT_TEMPLATE,
    input_variables=["conversation_history", "question"],
)

# ===============================
# PUBLIC API
# ===============================
@lru_cache(maxsize=1)
def get_llm():
    if LLM_PROVIDER == "groq":
        if not os.getenv("GROQ_API_KEY"):
            raise ValueError("GROQ_API_KEY is required when LLM_PROVIDER=groq.")
        return ChatGroq(
            model=GROQ_MODEL,
            temperature=0.0,
        )
    if LLM_PROVIDER == "gemini":
        if not os.getenv("GOOGLE_API_KEY"):
            raise ValueError("GOOGLE_API_KEY is required when LLM_PROVIDER=gemini.")
        return ChatGoogleGenerativeAI(
            model=GEMINI_MODEL,
            temperature=0.0,
            convert_system_message_to_human=True,
        )
    raise ValueError(
        f"Unsupported LLM_PROVIDER '{LLM_PROVIDER}'. Use 'groq' or 'gemini'."
    )


def get_rag_chain():
    llm = get_llm()
    retriever = BranhamRetriever()

    chain = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever,
        chain_type="stuff",
        return_source_documents=True,
        chain_type_kwargs={
            "prompt": PROMPT,
            "document_prompt": DOCUMENT_PROMPT,
            "document_variable_name": "context_str",
        },
        input_key="question",
    )

    return chain


def format_documents_for_context(docs: List[Document]) -> str:
    return "\n\n".join(
        DOCUMENT_PROMPT.format(
            page_content=doc.page_content,
            source=doc.metadata.get("source", ""),
            paragraph=doc.metadata.get("paragraph", ""),
        )
        for doc in docs
    )


def classify_grounding(answer: str, docs: List[Document]) -> str:
    normalized = answer.lower()
    if not docs or UNSUPPORTED_ANSWER.lower() in normalized:
        return "insufficient"
    if "not enough support" in normalized or "partial" in normalized or "not clearly supported" in normalized:
        return "partial"
    if len(docs) < 3:
        return "partial"
    return "supported"


def format_conversation_history(history: List[dict[str, str]], max_messages: int = 8) -> str:
    rows = []
    for item in history[-max_messages:]:
        role = item.get("role", "").strip().lower()
        content = " ".join(item.get("content", "").split())
        if role not in {"user", "assistant"} or not content:
            continue
        if len(content) > 900:
            content = content[:900].rstrip() + "..."
        rows.append(f"{role.upper()}: {content}")
    return "\n".join(rows) if rows else "No prior conversation."


def rewrite_conversational_question(question: str, history: List[dict[str, str]]) -> str:
    if not history:
        return question
    prompt_text = QUESTION_REWRITE_PROMPT.format(
        conversation_history=format_conversation_history(history),
        question=question,
    )
    response = get_llm().invoke(prompt_text)
    rewritten = getattr(response, "content", str(response)).strip()
    return rewritten or question


def normalize_answer_mode(answer_mode: str) -> str:
    if answer_mode not in ANSWER_MODE_INSTRUCTIONS:
        return "Strict Mode"
    return answer_mode


def build_answer_prompt(question: str, answer_mode: str, docs: List[Document]) -> str:
    context_str = format_documents_for_context(docs)
    return PROMPT_TEMPLATE.format(
        context_str=context_str,
        question=question,
        answer_mode_instructions=ANSWER_MODE_INSTRUCTIONS[answer_mode],
    )


def build_conversational_answer_prompt(
    question: str,
    retrieval_question: str,
    answer_mode: str,
    docs: List[Document],
    history: List[dict[str, str]],
) -> str:
    context_str = format_documents_for_context(docs)
    latest_question = question
    if retrieval_question != question:
        latest_question = f"{question}\n\nStandalone retrieval question: {retrieval_question}"
    return CONVERSATIONAL_PROMPT_TEMPLATE.format(
        context_str=context_str,
        conversation_history=format_conversation_history(history),
        question=latest_question,
        answer_mode_instructions=ANSWER_MODE_INSTRUCTIONS[answer_mode],
    )


def add_generation_diagnostics(
    diagnostics: dict[str, Any],
    total_started_at: float,
    answer_mode: str,
) -> dict[str, Any]:
    diagnostics["total_latency"] = round(time.perf_counter() - total_started_at, 3)
    diagnostics["answer_mode"] = answer_mode
    diagnostics["llm_provider"] = LLM_PROVIDER
    diagnostics["llm_model"] = GROQ_MODEL if LLM_PROVIDER == "groq" else GEMINI_MODEL
    return diagnostics


def build_refusal_result(
    diagnostics: dict[str, Any],
    total_started_at: float,
    answer_mode: str,
) -> dict[str, Any]:
    diagnostics["retrieved_source_count"] = diagnostics.get("source_count", 0)
    diagnostics["retrieved_distinct_sermon_count"] = diagnostics.get("distinct_sermon_count", 0)
    diagnostics["source_count"] = 0
    diagnostics["distinct_sermon_count"] = 0
    diagnostics["sermons"] = []
    add_generation_diagnostics(diagnostics, total_started_at, answer_mode)
    return {
        "answer": UNSUPPORTED_ANSWER,
        "source_documents": [],
        "grounding_status": "insufficient",
        "diagnostics": diagnostics,
    }


def prepare_answer_context(
    question: str,
    answer_mode: str = "Strict Mode",
    history: List[dict[str, str]] | None = None,
    conversational: bool = False,
) -> dict[str, Any]:
    total_started_at = time.perf_counter()
    answer_mode = normalize_answer_mode(answer_mode)
    history = history or []
    retrieval_question = (
        rewrite_conversational_question(question, history)
        if conversational and history
        else question
    )
    docs, diagnostics = retrieve_documents_with_diagnostics(
        retrieval_question,
        answer_mode=answer_mode,
    )
    diagnostics["retrieval_question"] = retrieval_question
    diagnostics["conversation_mode"] = conversational

    if diagnostics.get("should_refuse_generation"):
        return {
            "answer_mode": answer_mode,
            "docs": [],
            "diagnostics": diagnostics,
            "prompt_text": "",
            "total_started_at": total_started_at,
            "refusal_result": build_refusal_result(diagnostics, total_started_at, answer_mode),
        }

    return {
        "answer_mode": answer_mode,
        "docs": docs,
        "diagnostics": diagnostics,
        "prompt_text": (
            build_conversational_answer_prompt(
                question,
                retrieval_question,
                answer_mode,
                docs,
                history,
            )
            if conversational
            else build_answer_prompt(question, answer_mode, docs)
        ),
        "total_started_at": total_started_at,
        "refusal_result": None,
    }


def finalize_answer_result(
    answer: str,
    docs: List[Document],
    diagnostics: dict[str, Any],
    total_started_at: float,
    answer_mode: str,
) -> dict[str, Any]:
    grounding_status = classify_grounding(answer, docs)
    add_generation_diagnostics(diagnostics, total_started_at, answer_mode)
    return {
        "answer": answer,
        "source_documents": docs,
        "grounding_status": grounding_status,
        "diagnostics": diagnostics,
    }


def answer_question(
    question: str,
    answer_mode: str = "Strict Mode",
    history: List[dict[str, str]] | None = None,
    conversational: bool = False,
):
    prepared = prepare_answer_context(
        question,
        answer_mode,
        history,
        conversational,
    )
    if prepared["refusal_result"] is not None:
        return prepared["refusal_result"]

    llm = get_llm()
    response = llm.invoke(prepared["prompt_text"])
    answer = getattr(response, "content", str(response))
    return finalize_answer_result(
        answer,
        prepared["docs"],
        prepared["diagnostics"],
        prepared["total_started_at"],
        prepared["answer_mode"],
    )


def stream_answer_tokens(prompt_text: str) -> Iterator[str]:
    for chunk in get_llm().stream(prompt_text):
        text = getattr(chunk, "content", "")
        if isinstance(text, list):
            text = "".join(str(part) for part in text)
        if text:
            yield str(text)


def stream_answer_question(
    question: str,
    answer_mode: str = "Strict Mode",
    history: List[dict[str, str]] | None = None,
    conversational: bool = False,
) -> Iterator[dict[str, Any]]:
    prepared = prepare_answer_context(
        question,
        answer_mode,
        history,
        conversational,
    )
    if prepared["refusal_result"] is not None:
        result = prepared["refusal_result"]
        yield {
            "event": "metadata",
            "data": {
                "sources": result["source_documents"],
                "diagnostics": result["diagnostics"],
                "grounding_status": result["grounding_status"],
                "answer_mode": prepared["answer_mode"],
            },
        }
        yield {"event": "token", "data": {"text": result["answer"]}}
        yield {"event": "done", "data": result}
        return

    parts = []
    yield {
        "event": "metadata",
        "data": {
            "sources": prepared["docs"],
            "diagnostics": prepared["diagnostics"],
            "grounding_status": "pending",
            "answer_mode": prepared["answer_mode"],
        },
    }
    for token in stream_answer_tokens(prepared["prompt_text"]):
        parts.append(token)
        yield {"event": "token", "data": {"text": token}}

    answer = "".join(parts)
    result = finalize_answer_result(
        answer,
        prepared["docs"],
        prepared["diagnostics"],
        prepared["total_started_at"],
        prepared["answer_mode"],
    )
    yield {"event": "done", "data": result}


def search_archives(query: str):
    """
    Used by Search mode only.
    Returns (documents, debug_log)
    """
    debug = []
    docs = []
    seen = set()

    chunks = load_chunks()
    query_clean = normalize(query)

    # Keyword search
    for d in chunks:
        if query_clean in d.page_content.lower():
            key = d.page_content[:120]
            if key not in seen:
                docs.append(d)
                seen.add(key)

    debug.append(f"Keyword hits: {len(docs)}")

    # Fallback BM25
    if len(docs) < 20:
        for d in get_search_bm25_retriever().invoke(query):
            key = document_key(d)
            if key not in seen:
                docs.append(d)
                seen.add(key)

    debug.append(f"Total results: {len(docs)}")

    return docs, debug