# SPDX-FileCopyrightText: 2026 Surendra Singh <66422685+bolnet@users.noreply.github.com>
# SPDX-License-Identifier: MIT
"""HyDE retrieval — Hypothetical Document Embedding (Phase 3 PR-D).

A small LLM generates a 1-2-sentence hypothetical answer to the user's
question; the orchestrator embeds and searches BOTH the original
question AND the hypothetical answer in parallel, then RRF-merges the
two ranked lists before the rest of the cascade.

Per the RCA roadmap, this is the +6-10% recall lever sibling to
multi-query (PR #94). Intuition: dense embedding similarity between
a question and an answer is asymmetric — questions don't structurally
resemble their answers. Running both as queries widens the search
pattern across the latent space, raising the chance the correct
memory is in *some* lane's top-K.

This module is pure — it produces ``(queries: List[str], merged_hits:
List[dict])`` but does not call the vector store. The orchestrator
wires the lane in by handing in a ``vector_search`` callable, mirror
of the multi_query integration pattern.

Configuration lives in ``configs/attestor.yaml`` under
``retrieval.hyde``:

    retrieval:
      hyde:
        enabled: false           # default off
        generator_model: null    # null → models.extraction
        generator_reasoning_effort: low
        merge: rrf               # rrf | union

Trace events emitted (when ATTESTOR_TRACE=1):

  - ``recall.hyde.generated`` — the original + the hypothetical preview
  - ``recall.hyde.merged``    — pre-merge per-lane sizes + final size

Mutually exclusive with multi_query in this PR. The orchestrator
prefers multi_query when both are enabled (it's been shipped longer
and proven; HyDE is new). The combination is left for a follow-up.
"""

from __future__ import annotations

import asyncio
import logging
import os
from dataclasses import dataclass
from typing import Any
from collections.abc import Awaitable, Callable

logger = logging.getLogger("attestor.retrieval.hyde")


# Per-process AsyncOpenAI client cache, keyed by (base_url, api_key,
# timeout). The HyDE generator was previously building a fresh client
# per call, which leaks a TCP+TLS connection per question and pays the
# handshake cost (~200ms on cold cloud paths) on every request. Sharing
# a client lets the underlying httpx pool reuse connections across
# gathered HyDE lanes and across consecutive questions in a bench.
#
# Tests can pass ``client=`` directly to bypass the cache and inject a
# mock; production code uses the cache via the default arg.
_async_client_cache: dict[tuple[str | None, str, float], Any] = {}
# Lock the cache against concurrent first-callers under asyncio +
# to_thread fan-out; without this guard two threads can both miss
# the cache and build duplicate clients before the entry is populated.
_async_client_cache_lock = __import__("threading").Lock()


def _get_async_client(
    *, base_url: str | None, api_key: str, timeout: float,
) -> Any:
    """Return a process-wide cached ``AsyncOpenAI`` client for this
    (base_url, api_key, timeout) tuple. Construction is delegated to
    ``attestor.llm_trace.make_async_client`` so the trace wiring and
    explicit-timeout policy stay in one place.
    """
    key = (base_url, api_key, timeout)
    cached = _async_client_cache.get(key)
    if cached is not None:
        return cached
    with _async_client_cache_lock:
        cached = _async_client_cache.get(key)
        if cached is not None:
            return cached
        from attestor.llm_trace import make_async_client
        client = make_async_client(
            base_url=base_url,
            api_key=api_key,
            timeout=timeout,
        )
        _async_client_cache[key] = client
        return client


# ──────────────────────────────────────────────────────────────────────
# Result shape
# ──────────────────────────────────────────────────────────────────────


@dataclass(frozen=True)
class HydeResult:
    """The generator's output. ``original_question`` is always
    present; ``hypothetical_answer`` is empty string on degraded paths
    so the caller can detect "generation failed → just use the
    question alone" without checking for sentinel exceptions.
    """

    original_question: str
    hypothetical_answer: str = ""

    @property
    def queries(self) -> list[str]:
        """Queries to fan out, in lane-priority order. The original
        is always rank-0; the hypothetical comes second when present.

        On degraded paths (empty ``hypothetical_answer``), returns just
        the original — the caller naturally falls back to single-query.
        """
        if not self.hypothetical_answer.strip():
            return [self.original_question]
        return [self.original_question, self.hypothetical_answer]


# ──────────────────────────────────────────────────────────────────────
# Generator
# ──────────────────────────────────────────────────────────────────────


_GENERATOR_PROMPT = """You are previewing a personal-memory search. The \
user's question refers to an event, fact, or experience from their past \
conversations. Write 1-2 sentences DESCRIBING that event in the form the \
user would have ORIGINALLY MENTIONED it in conversation — not in the \
form of an answer to the question.

The goal is to generate a snippet whose embedding is close to the \
original conversation turn containing the answer, NOT to the answer \
itself. Question-shape and answer-shape don't embed close to source-shape.

Examples:
- Question: "How many weeks ago did I meet my aunt?"
  Snippet: "Yesterday I drove out to visit my aunt Linda — she gave me \
a gorgeous antique chandelier from her attic."
- Question: "Which trip did I take first this year?"
  Snippet: "Just got back from Tokyo last month, my first international \
trip of the year. The cherry blossoms were already starting."
- Question: "How many days ago did I buy a smoker?"
  Snippet: "Picked up a Traeger pellet smoker over the weekend — finally \
ready to do real low-and-slow brisket at home."

Rules:
- 1-2 sentences, first-person, declarative, conversational.
- DESCRIBE the event directly; do NOT say "X weeks ago" or "first" or \
counts — narrate the event itself.
- Invent plausible specifics (names, places, objects, dates) — accuracy \
doesn't matter, surface-form match to a likely chat turn does.
- Do NOT echo the question, do NOT hedge, do NOT refuse.

Question: {question}

Snippet:"""


def _resolve_generator_model() -> str:
    """Generator model: env override > YAML > extraction default."""
    if env := os.environ.get("HYDE_GENERATOR_MODEL"):
        return env
    from attestor.config import get_stack
    stack = get_stack()
    hyde_cfg = getattr(stack.retrieval, "hyde", None)
    if hyde_cfg is not None and getattr(hyde_cfg, "generator_model", None):
        return hyde_cfg.generator_model
    return stack.models.extraction


def generate_hypothetical_answer(
    question: str,
    *,
    model: str | None = None,
    api_key: str | None = None,
    timeout: float = 30.0,
) -> HydeResult:
    """Single LLM call returning a hypothetical answer for ``question``.

    On any error (missing key, malformed response, network timeout),
    returns ``HydeResult(original_question=question, hypothetical_answer="")``
    — the caller naturally degrades to a single-query lane.

    Provider URL + API key env are sourced from the YAML-configured LLM
    pool (``stack.llm.providers``). The ``api_key`` parameter is kept
    for back-compat: when set, it overrides the pool-resolved client by
    building a fresh one against the resolved provider's ``base_url``.
    """
    if not question.strip():
        return HydeResult(original_question=question.strip())

    model = model or _resolve_generator_model()

    try:
        from attestor.llm_trace import (
            get_client_for_model,
            _get_pool,
            make_client,
            traced_create,
        )
        if api_key:
            pool = _get_pool()
            head, sep, tail = model.partition("/")
            if sep and head in pool.providers:
                strategy = pool.get_strategy(head)
                clean_model = tail
            else:
                strategy = pool.default_strategy()
                clean_model = model
            client = make_client(
                base_url=strategy.base_url,
                api_key=api_key,
                timeout=timeout,
            )
        else:
            client, clean_model = get_client_for_model(model)
        response = traced_create(
            client,
            role="hyde_generator",
            model=clean_model,
            max_tokens=400,
            temperature=0.0,
            messages=[
                {"role": "user", "content": _GENERATOR_PROMPT.format(
                    question=question,
                )},
            ],
        )
        text = (response.choices[0].message.content or "").strip()
    except Exception as e:  # noqa: BLE001
        logger.debug("hyde.generate: LLM call failed: %s", e)
        return HydeResult(original_question=question.strip())

    # Strip any leading label the model may echo from the prompt frame.
    for label in ("snippet:", "hypothetical answer:"):
        if text.lower().startswith(label):
            text = text[len(label):].strip()
            break

    from attestor import trace as _tr
    if _tr.is_enabled():
        _tr.event(
            "recall.hyde.generated",
            original=question[:200],
            hypothetical_length=len(text),
            hypothetical_preview=text[:200],
        )

    return HydeResult(
        original_question=question.strip(),
        hypothetical_answer=text,
    )


# ──────────────────────────────────────────────────────────────────────
# Async generator (Phase 1 of the async retrieval rollout, see
# docs/plans/async-retrieval/PLAN.md). Mirrors generate_hypothetical_answer
# but awaits the LLM call so the caller can asyncio.gather it with the
# original-question vector embed instead of waiting serially.
#
# Audit invariant A7 (temperature=0 determinism) — same temperature=0.0
# pin as the sync path. Async amplifies non-determinism risk because
# gathered lanes can observe different hypotheticals across runs if
# T > 0; we explicitly pin T=0 here.
# ──────────────────────────────────────────────────────────────────────


async def generate_hypothetical_answer_async(
    question: str,
    *,
    model: str | None = None,
    api_key: str | None = None,
    timeout: float = 30.0,
    client: Any | None = None,
) -> HydeResult:
    """Async sibling of ``generate_hypothetical_answer``.

    Same prompt, same temperature=0, same fallback semantics. Uses
    ``openai.AsyncOpenAI`` so the call can run concurrently with the
    original-question vector embed via ``asyncio.gather`` in
    ``hyde_search_async``.

    The pool itself is sync-only, so we resolve the
    ``LLMProviderStrategy`` (base_url + api_key_env) from the YAML pool
    and build the ``AsyncOpenAI`` client locally — YAML still drives
    the URL/key. The optional ``api_key`` arg overrides the env-derived
    key against the same resolved ``base_url``.
    """
    if not question.strip():
        return HydeResult(original_question=question.strip())

    model = model or _resolve_generator_model()

    try:
        from attestor.llm_trace import _get_pool
        pool = _get_pool()
        head, sep, tail = model.partition("/")
        if sep and head in pool.providers:
            strategy = pool.get_strategy(head)
            clean_model = tail
        elif sep:
            logger.debug(
                "hyde.generate_async: unknown provider %r in model %r",
                head, model,
            )
            return HydeResult(original_question=question.strip())
        else:
            strategy = pool.default_strategy()
            clean_model = model

        key = api_key or os.environ.get(strategy.api_key_env)
        if not key:
            logger.debug(
                "hyde.generate_async: %s not set; returning degraded result",
                strategy.api_key_env,
            )
            return HydeResult(original_question=question.strip())

        # DI seam: tests pass a mock ``client`` directly; production
        # uses the process-wide cache so connections are pooled.
        active_client = client or _get_async_client(
            base_url=strategy.base_url,
            api_key=key,
            timeout=timeout,
        )
        response = await active_client.chat.completions.create(
            model=clean_model,
            max_tokens=400,
            temperature=0.0,
            messages=[
                {"role": "user", "content": _GENERATOR_PROMPT.format(
                    question=question,
                )},
            ],
        )
        text = (response.choices[0].message.content or "").strip()
    except Exception as e:  # noqa: BLE001
        logger.debug("hyde.generate_async: LLM call failed: %s", e)
        return HydeResult(original_question=question.strip())

    for label in ("snippet:", "hypothetical answer:"):
        if text.lower().startswith(label):
            text = text[len(label):].strip()
            break

    from attestor import trace as _tr
    if _tr.is_enabled():
        _tr.event(
            "recall.hyde.generated.async",
            original=question[:200],
            hypothetical_length=len(text),
            hypothetical_preview=text[:200],
        )

    return HydeResult(
        original_question=question.strip(),
        hypothetical_answer=text,
    )


# ──────────────────────────────────────────────────────────────────────
# End-to-end lane
# ──────────────────────────────────────────────────────────────────────


def hyde_search(
    question: str,
    *,
    vector_search: Callable[[str], list[dict[str, Any]]],
    model: str | None = None,
    merge: str = "rrf",
    api_key: str | None = None,
) -> tuple[list[str], list[dict[str, Any]]]:
    """Run the generator, fan out to ``vector_search`` for each query
    (original + hypothetical), merge the lanes, return
    ``(queries_used, merged_hits)``.

    Same pattern as ``attestor.retrieval.multi_query.multi_query_search``
    — RRF helpers are imported from there to avoid duplication. Falls
    back to single-query when generation fails (queries_used == [original]).
    """
    result = generate_hypothetical_answer(
        question, model=model, api_key=api_key,
    )
    queries = result.queries

    lanes: list[list[dict[str, Any]]] = []
    for q in queries:
        try:
            hits = list(vector_search(q))
        except Exception as e:  # noqa: BLE001
            logger.debug("hyde: lane %r failed: %s", q[:60], e)
            hits = []
        lanes.append(hits)

    # Reuse multi_query's merge helpers — same RRF k=60, same union
    # behavior. Single source of truth for ranked-list fusion.
    from attestor.retrieval.multi_query import (
        reciprocal_rank_fusion, union_merge,
    )
    if merge == "union":
        merged = union_merge(lanes)
    else:
        merged = reciprocal_rank_fusion(lanes)

    from attestor import trace as _tr
    if _tr.is_enabled():
        _tr.event(
            "recall.hyde.merged",
            n_queries=len(queries),
            per_lane_sizes=[len(l) for l in lanes],
            merged_size=len(merged),
            merge_strategy=merge,
        )

    return queries, merged


# ──────────────────────────────────────────────────────────────────────
# Async end-to-end lane (Phase 1 — docs/plans/async-retrieval/PLAN.md)
#
# Two concurrency levers:
#   1. Generator LLM call runs CONCURRENTLY with the original-question
#      vector search (asyncio.gather). The LLM is the slowest step;
#      overlapping it with the orig-question lane makes the orig lane
#      "free" on the wallclock.
#   2. Once the hypothetical is available, the hyde-lane vector search
#      runs (cannot start earlier — it needs the LLM output). Per-lane
#      failures are isolated via gather(..., return_exceptions=True);
#      RRF degrades to remaining lanes.
#
# Audit invariants preserved:
#   A7 — generator pins temperature=0.0 (same as sync).
#   The merged ranking is the deterministic RRF over both lanes given
#   identical inputs (verified in test_hyde_search_async_returns_same_result_as_sync).
# ──────────────────────────────────────────────────────────────────────


async def hyde_search_async(
    question: str,
    *,
    vector_search: Callable[[str], Awaitable[list[dict[str, Any]]]],
    model: str | None = None,
    merge: str = "rrf",
    api_key: str | None = None,
    timeout: float = 30.0,
) -> tuple[list[str], list[dict[str, Any]]]:
    """Async sibling of ``hyde_search``. ``vector_search`` MUST be an
    async callable returning a list of hit dicts.

    Latency story (with both LLM and vector search ~200ms each):
        sync:  ~600ms (generator → orig → hyde)
        async: ~400ms (generator ‖ orig, then hyde)

    On generator timeout or exception, degrades to single-lane recall
    (original-question only) without raising.
    """
    # Phase A: generator + orig-question vector search in parallel.
    # The orig task can run independent of the LLM — saves ~LLM latency.
    gen_task = asyncio.create_task(
        generate_hypothetical_answer_async(
            question, model=model, api_key=api_key, timeout=timeout,
        )
    )
    orig_task = asyncio.create_task(_safe_async_call(vector_search, question))

    # Wait for the generator with the same timeout. If it exceeds,
    # gracefully degrade to single-lane.
    try:
        result = await asyncio.wait_for(gen_task, timeout=timeout)
    except asyncio.TimeoutError:
        logger.debug("hyde_async: generator exceeded timeout=%.2fs", timeout)
        result = HydeResult(original_question=question.strip())
    except Exception as e:  # noqa: BLE001
        logger.debug("hyde_async: generator failed: %s", e)
        result = HydeResult(original_question=question.strip())

    queries = result.queries

    # Phase B: hyde-lane search (only if we got a hypothetical) and
    # await the original-lane task. Both use return_exceptions
    # semantics to avoid one lane's failure cancelling the other.
    if result.hypothetical_answer.strip():
        hyde_task = asyncio.create_task(
            _safe_async_call(vector_search, result.hypothetical_answer)
        )
        lanes_results = await asyncio.gather(
            orig_task, hyde_task, return_exceptions=True,
        )
    else:
        lanes_results = await asyncio.gather(orig_task, return_exceptions=True)

    # Filter exceptions / coerce to lists. Each successful lane is a
    # ``List[Dict]``. Exceptions become empty lanes (RRF handles).
    lanes: list[list[dict[str, Any]]] = []
    for lr in lanes_results:
        if isinstance(lr, Exception):
            lanes.append([])
        else:
            lanes.append(lr or [])

    # Reuse the sync RRF/union helpers — single source of truth for fusion.
    from attestor.retrieval.multi_query import (
        reciprocal_rank_fusion, union_merge,
    )
    if merge == "union":
        merged = union_merge(lanes)
    else:
        merged = reciprocal_rank_fusion(lanes)

    from attestor import trace as _tr
    if _tr.is_enabled():
        _tr.event(
            "recall.hyde.merged.async",
            n_queries=len(queries),
            per_lane_sizes=[len(l) for l in lanes],
            merged_size=len(merged),
            merge_strategy=merge,
        )

    return queries, merged


async def _safe_async_call(
    fn: Callable[[str], Awaitable[list[dict[str, Any]]]],
    arg: str,
) -> list[dict[str, Any]]:
    """Await ``fn(arg)``; on exception, return ``[]`` so the lane
    contributes an empty list to RRF rather than failing the whole
    recall. The exception is logged at DEBUG so audit traces capture
    "what we tried" even when a lane crashes.
    """
    try:
        return list(await fn(arg))
    except Exception as e:  # noqa: BLE001
        logger.debug("hyde_async: lane %r failed: %s", arg[:60], e)
        return []
