"""LLM-based relevance filter using plain text output.

Filters search previews by asking the LLM to return a list of relevant
indices as plain text (e.g. ``0, 2, 5``). We parse the response with a
regex over integers, which is robust to wrappers like "Indices: 0, 2"
or "[0, 2, 5]" and dodges all the structured-output provider quirks
(qwen prose-mode, function_calling latency, schema bikeshedding).

Design notes:
- An empty LLM response is treated as a valid judgment ("none of these
  results are relevant"). We do not second-guess the model — if the
  filter says reject all, we reject all, and log a warning so users
  can notice if their chosen model is misbehaving.
- On exception (network error, parse failure, provider outage) the
  filter is considered unavailable, not "reject all". In that case we
  fall back to a capped slice of the original previews so downstream
  processing is not overwhelmed by unfiltered results.
- The filter can split large preview lists into smaller ``batch_size``
  chunks. Smaller batches are faster per call and tend to be more
  reliable on weaker models which struggle to track many indices in a
  single context. A failed individual batch is skipped (logged); only
  a hard exception falls back to the capped slice.
- Partial success is the common outcome under batching: some batches
  return valid judgments, others raise or time out. Successful batches
  are kept; failed/timed-out batches contribute nothing. The capped
  fallback only fires when *every* batch failed to produce a result.
"""

import re
import time
from concurrent.futures import (
    ThreadPoolExecutor,
    TimeoutError as FuturesTimeoutError,
    as_completed,
)
from datetime import UTC, datetime
from typing import Any, Dict, List, Optional

from loguru import logger

from langchain_ollama import ChatOllama

from ..config.constants import DEFAULT_MAX_FILTERED_RESULTS

# Real wrapper chains (RateLimitedLLMWrapper -> ProcessingLLMWrapper -> base)
# are at most 2-3 levels deep. A depth limit avoids spurious infinite chains
# on test mocks where every getattr produces a new child object.
_MAX_UNWRAP_DEPTH = 10

# Matches standalone integers in the LLM response. ``\b`` ensures we
# don't match digits glued to letters ("v2"), and the negative lookbehind
# for ``.`` rejects the fractional part of decimals ("0.5" → just "0").
# Combined with the prompt instruction to output ONLY the numbers, this
# is robust against prose like "The top 3 results from 2024 are ...".
_INT_RE = re.compile(r"(?<![\w.])\d+\b")

# Wall-clock timeout passed to ``as_completed`` when batches run in
# parallel. Note: this bounds the time from iteration start until all
# parallel batches complete, not per-batch. A single batch (≤ 10
# previews on a 9B-class local model) typically completes in 5-15s.
# When Ollama serializes requests (OLLAMA_NUM_PARALLEL=1, the default)
# 10 parallel submissions effectively run sequentially server-side, so
# 10 × ~15s already brushes against the old 120s bound. 300s leaves
# room for slow-but-progressing work on serialized backends while still
# flagging genuine hangs. The sequential (workers==1) path does not
# apply this — only the parallel branch uses ``as_completed``.
# Without a bound here a stuck Ollama socket would block the pipeline
# indefinitely, because langchain-ollama's httpx client has no default
# socket timeout.
_FILTER_WALL_TIMEOUT_S = 300.0


# Per-preview snippet character cap. The snippet field often carries the
# paper abstract for academic engines (arxiv, semantic scholar, pubmed).
# 200 was too tight — it truncated abstracts before the judge could tell
# whether a paper's primary topic actually matches the query, leading to
# "paper mentions LLMs in passing" sources leaking past the filter. 800
# comfortably fits a typical abstract opening while keeping total prompt
# size bounded (≈16KB for 20 previews).
_SNIPPET_CHAR_CAP = 800


_RELEVANCE_PROMPT_TEMPLATE = """This is a relevance-filtering step. Kept results move forward and may be used in the final answer; dropped ones are excluded from further processing.

Query: "{query}"
Current date: {current_date}

Search results:
{preview_text}

Direct topic match matters more than keyword match — results that just mention the query terms usually don't help.

Output ONLY the 0-based indices of relevant results as a comma-separated list, nothing else.
Example: 0, 2, 5"""  # noqa: S608


def _unwrap_llm(llm):
    """Unwrap known LLM wrapper chains to get the base LangChain LLM.

    Walks ``.base_llm`` attributes up to ``_MAX_UNWRAP_DEPTH`` levels.
    The depth limit guards against test mocks (e.g. ``unittest.mock.Mock``)
    that lazily create a fresh child object on every attribute access.
    """
    probe = llm
    for _ in range(_MAX_UNWRAP_DEPTH):
        inner = getattr(probe, "base_llm", None)
        if inner is None or inner is probe:
            return probe
        probe = inner
    return probe


def _build_batch_prompt(
    query: str,
    batch: List[Dict[str, Any]],
    total_in_full: int,
    prompt_template: str,
) -> str:
    """Build the relevance prompt for a single batch of previews.

    Indices in the prompt are local to the batch (0..len(batch)-1).
    ``total_in_full`` is the size of the original full preview list,
    shown to the model for context — it doesn't affect the index range.
    ``prompt_template`` is passed in (defaulting at the public entry
    point) rather than read from a module global, so eval harnesses can
    test variants without monkey-patching module state.
    """
    preview_lines = []
    for i, preview in enumerate(batch):
        title = preview.get("title", "Untitled").strip()
        snippet = preview.get("snippet", "").strip()
        if len(snippet) > _SNIPPET_CHAR_CAP:
            snippet = snippet[:_SNIPPET_CHAR_CAP] + "..."
        preview_lines.append(f"[{i}] {title}\n    {snippet}")

    return prompt_template.format(
        query=query,
        current_date=datetime.now(UTC).strftime("%Y-%m-%d"),
        preview_text="\n\n".join(preview_lines),
    )


def _run_batch(
    llm,
    batch: List[Dict[str, Any]],
    query: str,
    total_in_full: int,
    engine_name: str,
    prompt_template: str,
) -> List[int]:
    """Invoke the LLM on a single batch and return the parsed local indices.

    Empty list = "none relevant" (valid judgment). Raises on LLM
    exceptions; the caller falls back to a capped slice in that case.
    """
    prompt = _build_batch_prompt(query, batch, total_in_full, prompt_template)
    return _invoke_text(llm, prompt, engine_name)


def filter_previews_for_relevance(
    llm,
    previews: List[Dict[str, Any]],
    query: str,
    max_filtered_results: Optional[int] = None,
    engine_name: str = "",
    batch_size: Optional[int] = None,
    max_parallel_batches: int = 1,
    prompt_template: Optional[str] = None,
) -> List[Dict[str, Any]]:
    """Filter search previews for relevance via plain-text LLM output.

    Args:
        llm: LangChain LLM instance (may be wrapped)
        previews: List of preview dicts with title/snippet/url
        query: The search query
        max_filtered_results: Optional cap on the final result count
            (None = LLM decides)
        engine_name: Engine class name for log messages
        batch_size: If set and smaller than ``len(previews)``, the LLM
            is called once per batch of this many previews. Smaller
            batches are faster per call and more reliable on weaker
            models. None or 0 disables batching (single-call mode).
        max_parallel_batches: Number of batches to dispatch concurrently
            against the LLM (via a thread pool). 1 = sequential.
            Most providers (Ollama with OLLAMA_NUM_PARALLEL>1, OpenAI,
            Anthropic) handle concurrent requests fine. Ignored when
            there is only one batch.
        prompt_template: Override the built-in ``_RELEVANCE_PROMPT_TEMPLATE``.
            Used by eval harnesses that want to compare variants without
            mutating module state. None (default) uses the shipping
            template.

    Returns:
        Filtered list of preview dicts (subset of input). Order matches
        the original preview order across batches.
    """
    if prompt_template is None:
        prompt_template = _RELEVANCE_PROMPT_TEMPLATE
    if not previews:
        return []

    for i, preview in enumerate(previews):
        title = preview.get("title", "Untitled").strip()
        logger.debug(f"[{engine_name}] INPUT [{i}]: {title[:80]}")

    # Cap used when the filter is unavailable (LLM exception) so we
    # don't flood downstream processing with unfiltered results.
    unavailable_cap = max_filtered_results or DEFAULT_MAX_FILTERED_RESULTS

    # Determine batch boundaries. A batch_size of None or 0 means
    # "single batch" — process all previews in one LLM call.
    effective_batch = (
        batch_size if (batch_size and batch_size > 0) else len(previews)
    )
    batch_starts = list(range(0, len(previews), effective_batch))
    batches = [previews[s : s + effective_batch] for s in batch_starts]

    workers = max(1, min(max_parallel_batches, len(batches)))
    logger.debug(
        f"[{engine_name}] Dispatching {len(batches)} batch(es) of "
        f"<= {effective_batch} previews each, {workers} parallel worker(s)"
    )

    t0 = time.monotonic()
    # None marks "batch never completed" (exception or timeout). A
    # completed batch is List[int], possibly empty (valid "none relevant"
    # judgment). We distinguish the two so an all-None outcome falls
    # back to the capped-slice while all-[] is kept as a valid decision.
    results_per_batch: List[Optional[List[int]]] = [None for _ in batches]

    if workers == 1:
        for i, batch in enumerate(batches):
            try:
                results_per_batch[i] = _run_batch(
                    llm,
                    batch,
                    query,
                    len(previews),
                    engine_name,
                    prompt_template,
                )
            except Exception:
                logger.exception(
                    f"[{engine_name}] batch {i} failed — skipping its results"
                )
    else:
        # Not using ``with ThreadPoolExecutor(...) as pool:`` on purpose.
        # ``__exit__`` calls ``shutdown(wait=True)`` which blocks on any
        # still-running worker — defeating the whole point of the
        # timeout when a worker is stuck on an Ollama HTTP call. We
        # manage the pool lifetime explicitly and always shut down with
        # ``wait=False`` so timed-out batches are orphaned (they'll
        # error out when the socket eventually fails) rather than
        # blocking the caller.
        pool = ThreadPoolExecutor(max_workers=workers)
        try:
            futures_to_idx = {
                pool.submit(
                    _run_batch,
                    llm,
                    batch,
                    query,
                    len(previews),
                    engine_name,
                    prompt_template,
                ): i
                for i, batch in enumerate(batches)
            }
            try:
                for fut in as_completed(
                    futures_to_idx, timeout=_FILTER_WALL_TIMEOUT_S
                ):
                    i = futures_to_idx[fut]
                    try:
                        results_per_batch[i] = fut.result()
                    except Exception:
                        logger.exception(
                            f"[{engine_name}] batch {i} failed — skipping"
                        )
            except FuturesTimeoutError:
                pending = [
                    futures_to_idx[f] for f in futures_to_idx if not f.done()
                ]
                logger.warning(
                    f"[{engine_name}] {len(pending)} batch(es) still "
                    f"running after {_FILTER_WALL_TIMEOUT_S}s — abandoning "
                    f"(pending indices: {pending})"
                )
        finally:
            # ``cancel_futures=True`` only cancels queued futures the
            # executor hasn't started yet; already-running workers stay
            # alive until their blocking call returns. We accept that —
            # the alternative is hanging forever.
            pool.shutdown(wait=False, cancel_futures=True)
    total_elapsed = time.monotonic() - t0

    # If no batch even completed, fall back to an unfiltered capped slice
    # — otherwise downstream would receive zero results from a filter
    # outage rather than from a valid "all irrelevant" judgment.
    if all(r is None for r in results_per_batch):
        logger.error(
            f"[{engine_name}] every relevance-filter batch failed — "
            f"returning first {unavailable_cap} previews as fallback"
        )
        return previews[:unavailable_cap]

    # Normalise None → [] for the aggregation pass below; a batch that
    # was cancelled or errored just contributes nothing. Assign to a
    # new variable with a narrower type so mypy can prove the list
    # elements are no longer Optional.
    completed_batches: List[List[int]] = [
        [] if r is None else r for r in results_per_batch
    ]

    # Aggregate results in original batch order so the final list
    # mirrors the input ordering across batches.
    ranked_results: List[Dict[str, Any]] = []
    kept_indices: List[int] = []
    seen: set = set()

    for batch_start, batch_result in zip(batch_starts, completed_batches):
        batch_len = min(effective_batch, len(previews) - batch_start)
        for li in batch_result:
            if not (0 <= li < batch_len):
                continue
            global_idx = batch_start + li
            if global_idx in seen:
                continue
            seen.add(global_idx)
            ranked_results.append(previews[global_idx])
            kept_indices.append(global_idx)

    logger.info(
        f"[{engine_name}] LLM relevance filter took {total_elapsed:.1f}s "
        f"across {len(batches)} batch(es) ({workers} parallel) "
        f"for {len(previews)} previews"
    )

    # Empty result is a valid LLM judgment ("none relevant"). Log a
    # warning on larger batches so users can notice a misbehaving model,
    # but do not override the decision.
    if not ranked_results and len(previews) > 2:
        logger.warning(
            f"[{engine_name}] LLM filter judged all {len(previews)} "
            f"results irrelevant. If this is unexpected, verify your "
            f"model handles structured output correctly."
        )

    # Apply cap if set, keeping ranked_results and kept_indices aligned.
    if (
        max_filtered_results is not None
        and len(ranked_results) > max_filtered_results
    ):
        ranked_results = ranked_results[:max_filtered_results]
        kept_indices = kept_indices[:max_filtered_results]

    # Log kept/removed/skipped. A preview is "skipped" when its batch
    # raised or timed out — the judge never saw it — which is distinct
    # from "removed" (judge returned a verdict that dropped it).
    skipped_indices: set = set()
    for i, (batch_start, r) in enumerate(zip(batch_starts, results_per_batch)):
        if r is None:
            batch_len = min(effective_batch, len(previews) - batch_start)
            skipped_indices.update(range(batch_start, batch_start + batch_len))
    removed_indices = (
        set(range(len(previews))) - set(kept_indices) - skipped_indices
    )
    logger.info(
        f"[{engine_name}] Relevance filter: "
        f"kept {len(ranked_results)} of {len(previews)} results"
    )
    for idx in kept_indices:
        title = previews[idx].get("title", "Untitled")[:80]
        logger.debug(f"[{engine_name}] KEPT    [{idx}]: {title}")
    for idx in sorted(removed_indices):
        title = previews[idx].get("title", "Untitled")[:80]
        logger.debug(f"[{engine_name}] REMOVED [{idx}]: {title}")
    for idx in sorted(skipped_indices):
        title = previews[idx].get("title", "Untitled")[:80]
        logger.debug(f"[{engine_name}] SKIPPED [{idx}]: {title}")

    return ranked_results


def _invoke_text(llm, prompt: str, engine_name: str) -> List[int]:
    """Invoke the LLM with a plain text prompt and parse out integer indices.

    Returns the list of parsed ints (empty list = "no integers found",
    treated as a valid "none relevant" judgment by the caller).
    Range/dedup validation happens in ``filter_previews_for_relevance``.
    """
    # Ollama thinking-by-default models (qwen3 dense variants, etc.)
    # burn 30-60s on CoT before emitting the answer. Index selection does
    # not benefit from reasoning, so suppress it on Ollama where supported.
    base_llm = _unwrap_llm(llm)
    invoke_kwargs = {}
    if isinstance(base_llm, ChatOllama):
        invoke_kwargs["reasoning"] = False

    result = llm.invoke(prompt, **invoke_kwargs)

    # LangChain chat models return a Message; LLMs return a string.
    text = getattr(result, "content", result)
    if not isinstance(text, str):
        logger.warning(
            f"[{engine_name}] Unexpected LLM response type: "
            f"{type(text).__name__}"
        )
        return []

    indices = [int(m) for m in _INT_RE.findall(text)]
    logger.debug(
        f"[{engine_name}] Text output parsed {len(indices)} indices: {indices}"
    )
    return indices
