app/services/rewriter.py

"""Author RAG Chatbot SaaS — Hybrid Query Rewriter.

3-Step Architecture (from RAG 1.2 proven pattern, adapted for book domain):
  Step 1 — Python follow-up resolver    → resolves pronouns + prepositional follows
  Step 2 — Python keyword expansion     → deterministic topic expansion map
  Step 3 — LLM rewrite (optional)       → only for short/vague/ambiguous queries

RULE: Python steps run first — zero API cost, microseconds.
RULE: LLM step is skipped if query is already clear after Step 1+2.
RULE: Always include original query in the returned list (first element).
"""

import re

import structlog

from app.config import get_settings

logger = structlog.get_logger(__name__)
cfg = get_settings()


# ── Step 1: Follow-Up Pronoun & Topic Resolver ───────────────────────────────
# Resolves "them", "it", "the book", "that" in follow-up queries using last
# user message from history. Used ONLY for retrieval — LLM gets original + history.
# Source: adapted from RAG 1.2 _resolve_followup()

_FOLLOWUP_PRONOUNS: set[str] = {
    "it", "them", "that", "this", "those", "these", "its",
    "the book", "the author", "the story", "the character",
}

_FOLLOWUP_STARTERS: tuple[str, ...] = (
    "and ", "also ", "what about ", "how about ",
    # Prepositional — almost always continue the previous topic
    "during ", "after ", "before ", "about ", "for ", "without ",
    "on ", "in ", "if ", "when ", "where ", "who ", "how ",
)


def _resolve_followup(query: str, history: list[dict]) -> str:
    """Resolve follow-up pronouns for better retrieval.

    Examples:
        "who is them?" + history["main characters"] → "who is main characters?"
        "and the villain?" → "main characters and the villain?"
        "how does it end?" → "how does [book title] end?"

    Only affects retrieval query — LLM receives original message + full history.

    Args:
        query: Current user message.
        history: Conversation history list.

    Returns:
        Resolved query string (may be same as input if no follow-up detected).
    """
    if not history:
        return query

    q_lower = query.lower().strip()
    words = set(q_lower.split())

    # Check for pronoun substitution
    has_pronoun = any(pron in q_lower for pron in _FOLLOWUP_PRONOUNS)

    # Check for short prepositional follow-up (e.g. "and the ending?")
    is_short_followup = (
        len(words) <= 6
        and any(q_lower.startswith(s) for s in _FOLLOWUP_STARTERS)
    )

    # Check for ultra-short vague query with no book topic keywords
    _BOOK_TOPIC_WORDS = {
        "buy", "purchase", "price", "author", "character", "plot", "story",
        "theme", "book", "chapter", "ending", "beginning", "setting", "genre",
        "review", "recommend", "summary", "spoil", "scene", "happen",
    }
    is_vague_short = len(words) <= 4 and not any(w in _BOOK_TOPIC_WORDS for w in words)

    if not has_pronoun and not is_short_followup and not is_vague_short:
        return query  # Nothing to resolve

    # Find the last user message for topic context
    last_user_msg = None
    for msg in reversed(history):
        if msg.get("role") == "user":
            last_user_msg = msg["content"].strip()
            break

    if not last_user_msg or last_user_msg.lower() == q_lower:
        return query  # No useful history

    # Replace first matching pronoun
    if has_pronoun:
        for pron in _FOLLOWUP_PRONOUNS:
            if pron in q_lower:
                resolved = re.sub(r'\b' + re.escape(pron) + r'\b', last_user_msg, query, count=1, flags=re.IGNORECASE)
                logger.debug("Follow-up pronoun resolved", original=query[:40], resolved=resolved[:40])
                return resolved

    # Prepend topic for short follow-ups
    if is_short_followup or is_vague_short:
        resolved = f"{last_user_msg} {query}"
        logger.debug("Follow-up prepended", original=query[:40], resolved=resolved[:40])
        return resolved

    return query


# ── Step 2: Book-Domain Keyword Expansion Map ─────────────────────────────────
# Deterministic keyword → search variant expansion. Runs in microseconds.
# Ordered most-specific → least-specific. First match only (break on hit).
# Source: original book-domain rules (adapted from RAG 1.2 _INTENT_MAP pattern)

_BOOK_EXPANSION_MAP: list[tuple[str, list[str]]] = [
    # ── Purchase / Buy (most specific first) ──────────────────────────────────
    ("how can i buy",    ["buy purchase where to order link available"]),
    ("where can i buy",  ["buy purchase where to order link available"]),
    ("where to buy",     ["buy purchase order link available store"]),
    ("how do i get",     ["buy purchase get copy order available"]),
    ("how to get",       ["buy purchase get copy order link"]),
    ("buy now",          ["buy purchase order link available"]),
    ("buy this",         ["buy purchase order link available"]),
    ("purchase",         ["buy purchase order link price available"]),
    ("how much",         ["price cost buy purchase link"]),
    ("price",            ["price cost buy purchase order link"]),
    ("buy",              ["buy purchase order link price available store"]),

    # ── Full Story / Spoiler ──────────────────────────────────────────────────
    ("tell me everything", ["full plot story summary complete overview spoiler ending"]),
    ("whole story",       ["complete story full plot summary all events"]),
    ("full story",        ["complete story full plot summary all events"]),
    ("complete summary",  ["summary plot overview story themes characters"]),
    ("how does it end",   ["ending conclusion resolution final plot"]),
    ("spoil",             ["spoiler ending conclusion plot resolution"]),
    ("summary",           ["summary overview plot story theme characters"]),

    # ── Characters ────────────────────────────────────────────────────────────
    ("who is the main",   ["main character protagonist hero role story"]),
    ("main character",    ["protagonist main character hero role"]),
    ("protagonist",       ["protagonist main character hero"]),
    ("villain",           ["antagonist villain evil character role"]),
    ("antagonist",        ["antagonist villain enemy character"]),
    ("who is",            ["character role person story protagonist antagonist"]),
    ("character",         ["character role person story protagonist antagonist"]),

    # ── Plot / Events ─────────────────────────────────────────────────────────
    ("what happen",       ["events plot story sequence what happens"]),
    ("what is it about",  ["main theme plot overview story topic subject"]),
    ("what is the book about", ["theme plot overview story topic summary"]),
    ("what happens",      ["plot events story sequence scene"]),
    ("plot",              ["plot story events narrative sequence theme"]),
    ("story",             ["story plot narrative events theme"]),

    # ── Themes / Message ──────────────────────────────────────────────────────
    ("what is the message", ["theme message moral lesson takeaway meaning"]),
    ("what does it mean",   ["theme meaning message moral lesson"]),
    ("moral",              ["moral lesson message theme takeaway"]),
    ("lesson",             ["lesson moral theme message takeaway"]),
    ("theme",              ["theme topic message moral meaning subject"]),
    ("message",            ["message theme moral lesson meaning"]),

    # ── Setting / World ───────────────────────────────────────────────────────
    ("where does",         ["setting location place world time period"]),
    ("when does",          ["time period era historical setting when"]),
    ("setting",            ["setting location time period world place"]),
    ("time period",        ["historical time period era when setting"]),

    # ── Recommendation / Audience ─────────────────────────────────────────────
    ("who should read",    ["target audience reader suitable recommend for"]),
    ("is it good for",     ["suitable recommend audience who should read"]),
    ("recommend",          ["recommend suitable audience reader who should read"]),
    ("who is it for",      ["target audience reader suitable recommend"]),
    ("good for",           ["suitable audience recommend reader who should"]),

    # ── Review / Quality ──────────────────────────────────────────────────────
    ("is it worth",        ["worth reading review quality good recommend"]),
    ("worth reading",      ["worth reading review quality recommend"]),
    ("is it good",         ["review quality worth reading recommend"]),
    ("review",             ["review quality opinion worth reading"]),

    # ── Author ───────────────────────────────────────────────────────────────
    ("who wrote",          ["author writer creator who wrote"]),
    ("author",             ["author writer creator wrote"]),

    # ── Genre / Style ────────────────────────────────────────────────────────
    ("genre",              ["genre category type fiction nonfiction style"]),
    ("what type",          ["genre category type fiction nonfiction style"]),
    ("what kind",          ["genre category type style audience"]),
    ("similar",            ["similar genre style comparable type"]),
    ("like this book",     ["similar genre style comparable"]),
]


def _expand_query(query: str) -> list[str]:
    """Expand query with domain-specific search variants. Deterministic, zero API cost.

    Inspired directly by RAG 1.2's _expand_query_locally(). On first match, breaks
    so only the most-specific expansion applies.

    Args:
        query: Query string (may already be follow-up resolved).

    Returns:
        List of query strings, original always first. Capped at 3.
    """
    q_lower = query.lower()
    queries: list[str] = [query]

    for keyword, variants in _BOOK_EXPANSION_MAP:
        if keyword in q_lower:
            for v in variants:
                if v not in queries:
                    queries.append(v)
            break  # Only first (most specific) match

    return queries[:3]


# ── Step 3: LLM Rewrite (ambiguous queries only) ──────────────────────────────
# Used for short/vague queries that Python cannot expand confidently.
# Returns plain text (not JSON) — simpler, faster, no parse failure risk.

_LLM_REWRITE_SYSTEM = """You are a book Q&A search query optimizer.

TASK: Rewrite the user query into clear, specific search terms for retrieving book content.
OUTPUT: A single line of search terms only — no explanation, no quotes, no prefix.

RULES:
- Resolve "it", "them", "that", "the book" using the conversation history
- Expand vague phrases into specific book-domain terms:
  * "tell me more" → "book content details story themes characters"
  * "what else?" → "additional book content story details"
  * "interesting" → "book themes story highlights notable moments"
- Keep output ≤20 words
- If query is already specific, output it unchanged"""


async def _llm_rewrite(query: str, history: list[dict]) -> str | None:
    """Step 3: LLM rewrite for genuinely ambiguous queries.

    Args:
        query: Current query (after Python steps).
        history: Conversation history.

    Returns:
        Rewritten query string or None on failure.
    """
    recent = history[-4:] if history else []
    history_str = "\n".join(
        f"{m['role'].title()}: {m['content'][:150]}"
        for m in recent
    ) or "None"

    user_msg = f"History:\n{history_str}\n\nQuery: {query}"

    try:
        # BUG-2 fix: use the shared singleton client instead of creating a new
        # AsyncOpenAI() here — was spawning a new HTTP connection pool on every call.
        from app.services.pipeline.helpers import _get_openai_client
        client = _get_openai_client()
        response = await client.chat.completions.create(
            model=cfg.OPENAI_CHAT_MODEL,
            messages=[
                {"role": "system", "content": _LLM_REWRITE_SYSTEM},
                {"role": "user", "content": user_msg},
            ],
            max_tokens=50,       # Plain text, very short
            temperature=0.1,
        )
        rewritten = response.choices[0].message.content.strip()
        # Strip any accidental quotes/prefixes from the model
        rewritten = re.sub(r'^["\']|["\']$', '', rewritten).strip()
        rewritten = re.sub(r'^(?:Rewritten|Search terms|Query)[:\s]+', '', rewritten, flags=re.IGNORECASE).strip()
        return rewritten if rewritten else None
    except Exception as e:
        logger.warning("LLM rewrite failed", error=str(e))
        return None


# ── Vague Query Detector ──────────────────────────────────────────────────────
# Determines if a query is too vague for Python expansion to help.

_VAGUE_PATTERNS: tuple[str, ...] = (
    "tell me more", "what else", "more about", "go on", "continue",
    "interesting", "cool", "nice", "wow", "really?", "and then?",
    "what do you think", "give me more", "elaborate", "explain more",
)


def _is_vague(query: str) -> bool:
    """Detect vague queries that benefit from LLM rewriting."""
    q = query.lower().strip()
    if len(q.split()) <= 3:
        return True
    return any(p in q for p in _VAGUE_PATTERNS)


# ── Public Interface ──────────────────────────────────────────────────────────

async def rewrite_query(query: str, history: list[dict]) -> list[str]:
    """Rewrite and expand a user query for maximum retrieval coverage.

    3-step hybrid approach:
      1. Python follow-up resolver (always runs, zero cost)
      2. Python keyword expansion (always runs, zero cost)
      3. LLM rewrite only if query is still vague (rarely needed)

    Args:
        query: Original user message text.
        history: Conversation history.

    Returns:
        List of query strings. Original always first. Maximum 4 total.
    """
    if not query.strip():
        return [query]

    # ── Step 1: Resolve follow-up pronouns (Python, zero cost) ───────────────
    resolved = _resolve_followup(query, history)

    # ── Step 2: Keyword expansion (Python, zero cost) ─────────────────────────
    queries = _expand_query(resolved)

    # Also expand the ORIGINAL query if it differs (the resolver may have mangled known terms)
    if resolved.lower() != query.lower():
        original_expansions = _expand_query(query)
        for oq in original_expansions:
            if oq not in queries:
                queries.append(oq)

    # Always ensure original is present
    if query not in queries:
        queries.insert(0, query)

    # ── Step 3: LLM rewrite for vague queries (rare API call) ─────────────────
    if _is_vague(resolved) and len(queries) <= 2:
        llm_rewritten = await _llm_rewrite(resolved, history)
        if llm_rewritten and llm_rewritten.lower() not in (q.lower() for q in queries):
            # Insert LLM rewrite as second item (after original, before expansions)
            queries.insert(1, llm_rewritten)
            logger.debug("LLM rewrite added", original=query[:40], rewritten=llm_rewritten[:40])

    logger.debug("Rewrite complete", original=query[:40], total=len(queries), queries=queries[:3])
    return queries[:4]  # Max 4: original + rewrite + 2 expansions