"""Cap the graded pool to N terms while keeping concept classes even.

The cap is a prefix of one fixed order: terms ranked by SHA-256 within their
concept class, merged round-robin across concepts. Evenness, reproducibility
and nesting all follow from that single order -- the sample at 40 terms is the
first 40 of the same sequence the sample at 80 takes from. Grader verdicts are
keyed on ``(term, field_of_law, question_id)``, so widening a batch must not
drop terms already graded.
"""

from __future__ import annotations

import hashlib
from collections import Counter
from collections.abc import Mapping
from dataclasses import dataclass


def _rank(term: str) -> int:
    return int(hashlib.sha256(term.encode("utf-8")).hexdigest(), 16)


@dataclass(frozen=True)
class TermSelection:
    """The capped term list together with the numbers behind it.

    Diagnostics render this rather than recounting the finished template, which
    has already been through the field-of-law assignment and no longer shows
    what the cap decided.
    """

    terms: tuple[str, ...]
    per_concept: Counter[str]
    available: Counter[str]
    requested: int | None

    def shortfall_concepts(self) -> list[str]:
        """Concepts that ran out before matching the fullest concept.

        These limited the sample; their unfilled slots went elsewhere. Empty
        whenever the whole pool was taken -- whether because there was no cap,
        or the cap met or exceeded the pool size -- since then nothing was
        left out and nothing is worth singling out.
        """
        if len(self.terms) >= sum(self.available.values()):
            return []
        fullest = max(self.per_concept.values(), default=0)
        return sorted(
            concept
            for concept, size in self.available.items()
            if size <= self.per_concept.get(concept, 0) < fullest
        )


def select_terms_by_concept(
    concept_by_term: Mapping[str, str],
    *,
    max_terms: int | None = None,
) -> TermSelection:
    """Take the first ``max_terms`` of the round-robin order over concepts."""
    pools: dict[str, list[str]] = {}
    # Sort the terms first so the rank sort is total even if two terms collide.
    for term in sorted(concept_by_term):
        pools.setdefault(concept_by_term[term], []).append(term)
    for terms in pools.values():
        terms.sort(key=_rank)

    concepts = sorted(pools)
    deepest = max((len(pool) for pool in pools.values()), default=0)
    order = [
        pools[concept][depth]
        for depth in range(deepest)
        for concept in concepts
        if depth < len(pools[concept])
    ]

    limit = len(order) if max_terms is None else max(0, max_terms)
    chosen = tuple(order[:limit])
    return TermSelection(
        terms=chosen,
        per_concept=Counter(concept_by_term[term] for term in chosen),
        available=Counter(concept_by_term.values()),
        requested=max_terms,
    )