"""
Stratified sampler for the L1 quality layer.

The L1 layer sends a small number of (source, translation) pairs to a
cheap LLM for a quality verdict. We don't want to send ALL chunks
(cost, latency) and we don't want to send random chunks (might miss
the problematic ones). We use a simple but effective strategy:

  - Prefer the LONGEST chunks (they contain the most diagnostic
    information per call).
  - Skip chunks that the L0 already flagged (we already know they're
    bad; we don't need the LLM to confirm).
  - Never sample more than `max_samples` chunks.
  - If `min_chunks` chunks aren't available, skip the L1 entirely
    (small documents don't need it).
"""

from __future__ import annotations

from typing import List, Tuple


def sample_chunks_for_l1(
    source_chunks: List[str],
    translated_chunks: List[str],
    failed_indices: set,
    max_samples: int = 5,
    min_chunks: int = 10,
) -> List[Tuple[str, str]]:
    """
    Select a sample of (source, translation) pairs to send to the L1 judge.

    Args:
        source_chunks: Original texts.
        translated_chunks: Translated texts (same length as source_chunks).
        failed_indices: Set of indices that L0 already flagged as bad.
                        These are SKIPPED — we want fresh signal, not
                        confirmation of an already-detected failure.
        max_samples: Maximum number of pairs to return.
        min_chunks: If the document has fewer than this many chunks,
                    return an empty list (not enough signal to bother
                    calling the LLM).

    Returns:
        A list of (source, translated) tuples, ready to send to the LLM
        judge. Order: longest chunks first.
    """
    n = min(len(source_chunks), len(translated_chunks))

    if n < min_chunks:
        return []

    # Build candidate list, excluding L0 failures
    candidates: List[Tuple[int, int, str, str]] = []
    for i in range(n):
        if i in failed_indices:
            continue
        src = (source_chunks[i] or "").strip()
        trans = (translated_chunks[i] or "").strip()
        # Skip very short pairs (not enough context for the LLM to judge)
        if len(src) < 5 and len(trans) < 5:
            continue
        # Skip pairs where source and translation are identical
        # (probably a non-translated cell like a number, code, or brand)
        if src == trans:
            continue
        # Rank by length of the translation (longer = more diagnostic)
        rank = len(trans) + len(src) // 2
        candidates.append((rank, i, src, trans))

    # Sort by length descending and take top N
    candidates.sort(key=lambda c: c[0], reverse=True)
    selected = candidates[:max_samples]

    # Return in original document order (helps the LLM judge maintain
    # context, and makes the verdict easier to interpret)
    selected.sort(key=lambda c: c[1])

    return [(src, trans) for _rank, _i, src, trans in selected]