"""LLM-based reranker — blends LLM relevance with retrieval scores.

The reranker asks an LLM to rate how relevant each candidate fact is to
the query. The LLM score is blended with the existing formula score via
a weighted average (configurable via ``reranker_weight``). This means the
reranker influences ranking but cannot single-handedly zero out a fact
that has strong semantic/keyword/graph signals.

Graceful degradation: if the LLM call fails, returns original ranking unchanged.
"""

import asyncio
import json
import logging

from arandu._helpers import strip_markdown_fences
from arandu.config import MemoryConfig
from arandu.protocols import LLMProvider
from arandu.read.retrieval import RetrievalCandidate

logger = logging.getLogger(__name__)

_RERANKER_SYSTEM_PROMPT = """\
You are a relevance judge for a long-term memory system.

You will receive a user query and a list of memory facts stored about that user.
Each fact is a piece of information previously extracted from the user's conversations.

Your task: rate how relevant each fact is to answering the query.

Scoring guide:
- 1.0: Directly answers the query or is essential context.
- 0.7-0.9: Highly relevant, provides important supporting information.
- 0.4-0.6: Somewhat relevant, tangentially related.
- 0.1-0.3: Barely relevant, only loosely connected.
- 0.0: Completely irrelevant to the query.

Important:
- A fact mentioning the same entity as the query is NOT automatically relevant.
  Judge by whether the fact helps ANSWER the query.
- A fact that contains the exact information being asked about should score 1.0,
  even if the entity names are unfamiliar to you.
- Return scores as a JSON object: {"scores": [score1, score2, ...]}
- The array must have exactly N elements, one per fact."""


async def rerank_with_llm(
    query: str,
    candidates: list[RetrievalCandidate],
    llm: LLMProvider,
    config: MemoryConfig,
) -> list[RetrievalCandidate]:
    """Rerank candidates using multiplicative blend of LLM relevance + formula score.

    The final score is:
        final = formula_score * (floor + w * reranker_score)

    where w = config.reranker_weight (default 0.70), floor = 1 - w = 0.30.

    reranker=1.0 → multiplier=1.0 (formula unchanged)
    reranker=0.0 → multiplier=0.30 (formula reduced to 30%)

    This means the reranker penalizes proportionally: strong formula facts
    survive even with reranker=0, weak formula facts get filtered.

    Falls back to original ranking on any failure.
    """
    if not candidates or len(candidates) <= 1:
        return candidates

    # Select pool for reranking (top candidates only)
    pool_size = min(len(candidates), config.topk_facts)
    pool = candidates[:pool_size]

    # Build fact display strings for the prompt
    fact_lines = []
    for i, c in enumerate(pool):
        f = c.fact
        display = f"{f.entity_name or ''}: {f.fact_text or ''}"
        fact_lines.append(f"{i + 1}. {display}")

    facts_block = "\n".join(fact_lines)
    user_prompt = f'Query: "{query}"\n\nFacts ({len(pool)} total):\n{facts_block}'

    try:
        _llm_result = await asyncio.wait_for(
            llm.complete(
                messages=[
                    {"role": "system", "content": _RERANKER_SYSTEM_PROMPT},
                    {"role": "user", "content": user_prompt},
                ],
                temperature=0,
                response_format={"type": "json_object"},
                max_tokens=300,
            ),
            timeout=config.reranker_timeout_sec,
        )
        raw = _llm_result.text

        raw = strip_markdown_fences(raw)
        data = json.loads(raw)
        scores = data.get("scores", [])

        if not isinstance(scores, list) or len(scores) != len(pool):
            logger.warning(
                "reranker: unexpected scores length=%s, expected=%s, falling back",
                len(scores) if isinstance(scores, list) else "N/A",
                len(pool),
            )
            return candidates

        # Multiplicative blend: final = formula * (floor + (1-floor) * reranker)
        # floor = 1 - reranker_weight. With default weight=0.70, floor=0.30.
        # reranker=1.0 → multiplier=1.0 (no change). reranker=0.0 → multiplier=0.30.
        w = max(0.0, min(1.0, config.reranker_weight))
        floor = 1.0 - w
        for i, c in enumerate(pool):
            try:
                reranker_score = max(0.0, min(1.0, float(scores[i])))
                c.scores["reranker"] = reranker_score
                formula_score = c.final_score  # pre-reranker score from retrieval
                c.scores["formula"] = formula_score  # preserve for debugging
                c.final_score = formula_score * (floor + w * reranker_score)
            except (TypeError, ValueError):
                c.scores["reranker"] = 0.0
                # Keep existing final_score on parse error

        # Sort pool by blended score descending
        reranked_pool = sorted(pool, key=lambda x: x.final_score, reverse=True)

        # Build final list: reranked pool + remaining candidates not in pool
        pool_ids = {id(c) for c in pool}
        rest = [c for c in candidates if id(c) not in pool_ids]
        return reranked_pool + rest

    except TimeoutError:
        logger.warning(
            "reranker: timeout after %.1fs, falling back to first-pass ranking",
            config.reranker_timeout_sec,
        )
        return candidates
    except Exception as e:
        logger.warning("reranker: failed: %s, falling back to first-pass ranking", e)
        return candidates
