"""Evaluation rubrics for each CRI dimension.

Each ``*_check`` function is a pure function with no side effects that accepts
structured inputs and returns a complete prompt string ready to be sent to an
LLM judge.

Most rubrics produce binary YES/NO prompts (consumed by
:meth:`~cri.judge.BinaryJudge.judge`).  The exception is
:func:`mei_coverage_chunk_check`, which produces a multi-answer prompt
consumed by :meth:`~cri.judge.BinaryJudge.judge_coverage` — the judge
returns a JSON array of covered GT-fact indices rather than YES/NO.

Design principles:
- MAX_FACTS_PER_PROMPT caps stored facts per prompt (context budget).
- ``format_facts`` provides consistent numbered-list formatting.
- Prompts emphasise **semantic equivalence** — meaning match, not exact text.
- For "negative" checks (staleness, noise, irrelevance) YES means the system
  **failed** the check.  The prompt makes this explicit so the caller can
  interpret verdicts uniformly.
"""

from __future__ import annotations

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

MAX_FACTS_PER_PROMPT: int = 30
"""Maximum number of stored facts included in a single judge prompt."""

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def format_facts(facts: list[str]) -> str:
    """Format a list of fact strings into a numbered list for prompts.

    If the list exceeds :data:`MAX_FACTS_PER_PROMPT` entries, it is truncated
    and a note is appended indicating how many facts were omitted.

    Args:
        facts: Raw fact strings to format.

    Returns:
        A formatted string with numbered facts, or ``"(no facts provided)"``
        if the list is empty.
    """
    if not facts:
        return "(no facts provided)"

    truncated = facts[:MAX_FACTS_PER_PROMPT]
    lines = [f"  {i}. {fact}" for i, fact in enumerate(truncated, start=1)]
    result = "\n".join(lines)

    remaining = len(facts) - MAX_FACTS_PER_PROMPT
    if remaining > 0:
        result += f"\n  [... {remaining} more facts not shown]"

    return result


# ---------------------------------------------------------------------------
# Binary verdict rubric functions
# ---------------------------------------------------------------------------


def pas_check(dimension: str, gold_answer: str, stored_facts: list[str]) -> str:
    """Generate a prompt to check if stored facts match a profile dimension's gold answer.

    Args:
        dimension: Name of the profile dimension (e.g. ``"occupation"``).
        gold_answer: The expected / ground-truth value for that dimension.
        stored_facts: Facts stored by the memory system under evaluation.

    Returns:
        A prompt string for the LLM judge (expected answer: YES or NO).
    """
    facts_block = format_facts(stored_facts)
    return (
        "TASK\n"
        "You are evaluating whether an AI memory system correctly captured a user's "
        "profile information. Determine if the stored facts contain information that "
        "semantically matches the expected value for the given profile dimension.\n"
        "Consider semantic equivalence: the stored fact does not need to use the exact "
        "same words — if the meaning is the same, that counts as a match.\n\n"
        f"PROFILE DIMENSION: {dimension}\n"
        f"EXPECTED VALUE: {gold_answer}\n\n"
        "STORED FACTS:\n"
        f"{facts_block}\n\n"
        "QUESTION\n"
        "Do the stored facts contain information that semantically matches the "
        f'expected value "{gold_answer}" for the profile dimension "{dimension}"?\n\n'
        "Answer YES or NO."
    )


def dbu_recency_check(fact_name: str, new_value: str, stored_facts: list[str]) -> str:
    """Generate a prompt to check if the current/updated value is reflected.

    Args:
        fact_name: Human-readable name of the fact (e.g. ``"job title"``).
        new_value: The updated / most-recent value.
        stored_facts: Facts stored by the memory system under evaluation.

    Returns:
        A prompt string for the LLM judge (expected answer: YES or NO).
    """
    facts_block = format_facts(stored_facts)
    return (
        "TASK\n"
        "You are evaluating whether an AI memory system has correctly updated its "
        "knowledge when new information was provided. Determine if the stored facts "
        "reflect the most recent / current value for the given fact.\n"
        "Consider semantic equivalence: the wording does not need to be identical — "
        "the meaning must match.\n\n"
        f"FACT: {fact_name}\n"
        f"EXPECTED CURRENT VALUE: {new_value}\n\n"
        "STORED FACTS:\n"
        f"{facts_block}\n\n"
        "QUESTION\n"
        "Do the stored facts reflect that the current value of "
        f'"{fact_name}" is "{new_value}" (or something semantically equivalent)?\n\n'
        "Answer YES or NO."
    )


def dbu_staleness_check(fact_name: str, old_value: str, stored_facts: list[str]) -> str:
    """Generate a prompt to check if an old value is still asserted as *current*.

    Historical mentions (e.g. "used to be X") are acceptable — this check only
    flags the old value if it is presented as the **current** truth.

    **Interpretation**: YES = the system still asserts the old value as current
    (this is a failure).  NO = the old value is not asserted as current (good).

    Args:
        fact_name: Human-readable name of the fact.
        old_value: The outdated value that should have been superseded.
        stored_facts: Facts stored by the memory system under evaluation.

    Returns:
        A prompt string for the LLM judge (expected answer: YES or NO).
    """
    facts_block = format_facts(stored_facts)
    return (
        "TASK\n"
        "You are evaluating whether an AI memory system still incorrectly asserts "
        "an outdated value as the CURRENT truth. The old value may appear in a "
        "historical context (e.g. 'previously was X', 'used to be X') — that is "
        "acceptable. It is only a problem if the old value is presented as the "
        "current, active value.\n"
        "Consider semantic equivalence when comparing values.\n\n"
        f"FACT: {fact_name}\n"
        f"OLD (OUTDATED) VALUE: {old_value}\n\n"
        "STORED FACTS:\n"
        f"{facts_block}\n\n"
        "QUESTION\n"
        "Do the stored facts still assert that the current value of "
        f'"{fact_name}" is "{old_value}" (or something semantically equivalent), '
        "treating it as the present truth rather than a historical note?\n"
        "NOTE: YES means the system FAILED to update — the old value is still "
        "treated as current.\n\n"
        "Answer YES or NO."
    )


def tc_temporal_validity_check(
    fact_description: str,
    expected_current: bool,
    stored_facts: list[str],
) -> str:
    """Generate a prompt to check temporal fact handling.

    When *expected_current* is ``True`` the judge checks that the fact IS
    present and treated as current.  When ``False`` the judge checks that the
    fact is NOT asserted as currently valid.

    Args:
        fact_description: Human-readable description of the temporal fact.
        expected_current: Whether the fact should be treated as currently valid.
        stored_facts: Facts stored by the memory system under evaluation.

    Returns:
        A prompt string for the LLM judge (expected answer: YES or NO).
    """
    facts_block = format_facts(stored_facts)

    if expected_current:
        question = "Is the fact described above present in the stored facts and treated as currently valid?"
        note = ""
    else:
        question = "Do the stored facts still assert the fact described above as currently valid, even though it should no longer be current?"
        note = "NOTE: YES means the system FAILED — it treats an expired / no-longer-current fact as still valid.\n"

    return (
        "TASK\n"
        "You are evaluating whether an AI memory system correctly handles "
        "the temporal validity of facts. Some facts are only valid for a "
        "certain period — the system should know which facts are current "
        "and which have expired or been superseded.\n"
        "Consider semantic equivalence when comparing.\n\n"
        f"FACT: {fact_description}\n"
        f"EXPECTED TO BE CURRENT: {'Yes' if expected_current else 'No'}\n\n"
        "STORED FACTS:\n"
        f"{facts_block}\n\n"
        "QUESTION\n"
        f"{question}\n"
        f"{note}\n"
        "Answer YES or NO."
    )


def crq_resolution_check(
    topic: str,
    correct_resolution: str,
    stored_facts: list[str],
) -> str:
    """Generate a prompt to check if a conflict was resolved correctly.

    Args:
        topic: The topic area where the conflict occurred.
        correct_resolution: The expected correct resolution.
        stored_facts: Facts stored by the memory system under evaluation.

    Returns:
        A prompt string for the LLM judge (expected answer: YES or NO).
    """
    facts_block = format_facts(stored_facts)
    return (
        "TASK\n"
        "You are evaluating whether an AI memory system correctly resolved "
        "a conflict between contradictory pieces of information. The system "
        "should have identified the conflict and arrived at the correct "
        "resolution.\n"
        "Consider semantic equivalence: the stored resolution does not need "
        "to use the exact same words — the meaning must match.\n\n"
        f"CONFLICT TOPIC: {topic}\n"
        f"EXPECTED CORRECT RESOLUTION: {correct_resolution}\n\n"
        "STORED FACTS:\n"
        f"{facts_block}\n\n"
        "QUESTION\n"
        "Do the stored facts reflect the correct resolution of the conflict "
        f'on "{topic}", consistent with the expected resolution '
        f'"{correct_resolution}"?\n\n'
        "Answer YES or NO."
    )


def qrp_relevance_check(
    query: str,
    expected_fact: str,
    stored_facts: list[str],
) -> str:
    """Generate a prompt to check if a relevant fact is present in results.

    Args:
        query: The query posed to the memory system.
        expected_fact: The fact expected to be present in the results.
        stored_facts: Facts returned by the memory system for the query.

    Returns:
        A prompt string for the LLM judge (expected answer: YES or NO).
    """
    facts_block = format_facts(stored_facts)
    return (
        "TASK\n"
        "You are evaluating the precision and relevance of facts returned by "
        "an AI memory system in response to a query. Determine whether the "
        "returned facts include information that matches the expected fact.\n"
        "Consider semantic equivalence: meaning match counts, not exact wording.\n\n"
        f"QUERY: {query}\n"
        f"EXPECTED FACT (should be present): {expected_fact}\n\n"
        "RETURNED FACTS:\n"
        f"{facts_block}\n\n"
        "QUESTION\n"
        "Do the returned facts contain information that semantically matches "
        f'the expected fact "{expected_fact}"?\n\n'
        "Answer YES or NO."
    )


def mei_coverage_check(
    gt_key: str,
    gt_value: str,
    stored_facts: list[str],
) -> str:
    """Generate a prompt to check if a ground-truth fact is covered.

    Args:
        gt_key: Name of the ground-truth dimension (e.g. ``"occupation"``).
        gt_value: The expected value for that dimension.
        stored_facts: All facts stored by the memory system.

    Returns:
        A prompt string for the LLM judge (expected answer: YES or NO).
    """
    facts_block = format_facts(stored_facts)
    return (
        "TASK\n"
        "You are evaluating whether an AI memory system's stored facts "
        "cover a specific expected ground-truth fact. The stored fact does "
        "not need to use the exact same words — semantic equivalence "
        "counts.\n\n"
        f"EXPECTED FACT\n"
        f"  Dimension: {gt_key}\n"
        f"  Value: {gt_value}\n\n"
        "STORED FACTS:\n"
        f"{facts_block}\n\n"
        "QUESTION\n"
        "Do ANY of the stored facts above convey the same information as "
        f'the expected fact "{gt_key}: {gt_value}" '
        "(or something semantically equivalent)?\n\n"
        "Answer YES or NO."
    )


def mei_coverage_chunk_check(
    chunk_facts: list[str],
    gt_facts: list[tuple[str, str]],
) -> str:
    """Generate a multi-answer prompt asking which GT facts a chunk covers.

    Used by the MEI dimension's chunk-outer scanning loop.  For each chunk of
    ``MAX_FACTS_PER_PROMPT`` stored facts the judge is asked to identify which
    ground-truth facts are represented — returning a JSON array of 0-based
    indices rather than a binary YES/NO verdict.

    Args:
        chunk_facts: A slice of stored facts (at most MAX_FACTS_PER_PROMPT).
        gt_facts: All ground-truth ``(dimension_name, value)`` pairs for the
            full evaluation run.  Indices here must stay stable across calls.

    Returns:
        A prompt string for :meth:`~cri.judge.BinaryJudge.judge_coverage`.
        The judge should respond with a JSON array of the 0-based indices of
        GT facts found in this chunk (e.g. ``[0, 2, 5]`` or ``[]``).
    """
    facts_block = format_facts(chunk_facts)
    gt_lines = "\n".join(f"  [{i}] {key}: {value}" for i, (key, value) in enumerate(gt_facts))
    return (
        "TASK\n"
        "You are evaluating whether a chunk of an AI memory system's stored facts "
        "covers any of the expected ground-truth facts listed below. For each "
        "ground-truth fact, determine if ANY stored fact in this chunk conveys the "
        "same information. Semantic equivalence counts — exact wording is not required.\n\n"
        "STORED FACTS (this chunk):\n"
        f"{facts_block}\n\n"
        "GROUND-TRUTH FACTS TO CHECK:\n"
        f"{gt_lines}\n\n"
        "INSTRUCTIONS\n"
        "Return ONLY a JSON array of the 0-based indices of the ground-truth facts "
        "that are covered by the stored facts above. "
        "If none are covered, return []. "
        "Do not include any explanation — only the JSON array.\n\n"
        "Example: [0, 2, 5]"
    )


def qrp_irrelevance_check(
    query: str,
    irrelevant_fact: str,
    stored_facts: list[str],
) -> str:
    """Generate a prompt to check if an irrelevant fact was incorrectly included.

    **Interpretation**: YES = irrelevant fact was included (failure).
    NO = irrelevant fact was correctly excluded (good).

    Args:
        query: The query posed to the memory system.
        irrelevant_fact: The fact that should NOT appear in the results.
        stored_facts: Facts returned by the memory system for the query.

    Returns:
        A prompt string for the LLM judge (expected answer: YES or NO).
    """
    facts_block = format_facts(stored_facts)
    return (
        "TASK\n"
        "You are evaluating whether an AI memory system incorrectly included "
        "an irrelevant fact in its response to a query. The system should "
        "return only facts relevant to the query.\n"
        "Consider semantic equivalence when comparing.\n\n"
        f"QUERY: {query}\n"
        f"IRRELEVANT FACT (should NOT be in results): {irrelevant_fact}\n\n"
        "RETURNED FACTS:\n"
        f"{facts_block}\n\n"
        "QUESTION\n"
        "Do the returned facts contain information that is essentially the "
        f'same as the irrelevant fact "{irrelevant_fact}"?\n'
        "NOTE: YES means the system FAILED — it included an irrelevant fact "
        "in its response.\n\n"
        "Answer YES or NO."
    )
