"""Bounded review batching after complete-suite discovery and reconciliation."""
from suite_contract import Problem
from suite_policy import MAX_GROUP, invocation, validate_natural

MAX_REVIEW_BATCHES = 8
MAX_MODEL_CALLS = 64


def material(stage, original, reviewed=None):
    """Retain every original oversized family for semantic review and its audit."""
    context = {"oversized_families": [g["members"] for g in original["groups"] if len(g["members"]) > MAX_GROUP]}
    if stage == "large_family_review":
        context["natural_partition"] = original
    else:
        context.update(original_partition=original, reviewed_partition=reviewed)
    return context


def batch_request(stage, source, originals, reviewed, index, total):
    """Slice only reconciled family boundaries, retaining all their source fields."""
    aliases = {a for g in originals for a in g["members"]}
    subset = {**source, "cases": [c for c in source["cases"] if c["alias"] in aliases]}
    prior = None if reviewed is None else {
        "groups": [g for g in reviewed["groups"] if set(g["members"]) <= aliases],
        "decisions": [d for d in reviewed["decisions"] if set(d["source_members"]) <= aliases]}
    context = material(stage, {"groups": originals}, prior)
    context["review_batch"] = {"index": index, "total": total,
                              "case_count": len(aliases), "whole_suite_case_count": len(source["cases"])}
    count = len(prior["groups"]) if prior else len(originals)
    return subset, context, count


def review_stage(workflow, stage, source, original, reviewed, capacity):
    """Use one full review when it fits, otherwise bounded whole-family batches.

    Discovery and reconciliation already considered every suite case together.
    Reviews may refine only within an original family under the existing policy;
    their execution batches therefore introduce no new semantic boundaries.
    """
    context = material(stage, original, reviewed)
    count = len((reviewed or original)["groups"])
    try:
        result = workflow.call(stage, source, context, count)
    except Problem as exc:
        if exc.code not in {"pass_capacity", "pass_request_too_large"}:
            raise
    else:
        validate_natural(result, source, original["groups"])
        return result

    large = sorted((g for g in original["groups"] if len(g["members"]) > MAX_GROUP),
                   key=lambda g: tuple(sorted(g["members"])))
    batches, current = [], []
    for family in large:
        candidate = current + [family]
        subset, context, count = batch_request(stage, source, candidate, reviewed, 8, 8)
        try:
            capacity(invocation(stage, subset, context), workflow.provider, len(subset["cases"]), count)
        except Problem as exc:
            if exc.code not in {"pass_capacity", "pass_request_too_large"}:
                raise
            if not current:
                raise Problem("review_family_capacity", 422, review_pass=stage,
                              case_count=len(family["members"]), failure_stage="review_batch_preflight") from None
            batches.append(current)
            current = [family]
        else:
            current = candidate
    if current:
        batches.append(current)
    if not batches or len(batches) > MAX_REVIEW_BATCHES:
        raise Problem("review_batch_limit", 422, review_pass=stage, batch_count=len(batches))
    calls = [batch_request(stage, source, group, reviewed, i, len(batches))
             for i, group in enumerate(batches, 1)]
    # Admit every batch before launching any paid review call.
    for subset, context, count in calls:
        capacity(invocation(stage, subset, context), workflow.provider, len(subset["cases"]), count)
    workflow.max_calls = MAX_MODEL_CALLS
    if workflow.max_calls > MAX_MODEL_CALLS:
        raise Problem("model_pass_limit", 502)
    combined = {"groups": [g for g in original["groups"] if len(g["members"]) <= MAX_GROUP], "decisions": []}
    base_reserve = workflow.future_seconds
    for index, (originals, (subset, context, count)) in enumerate(zip(batches, calls)):
        workflow.future_seconds = base_reserve + (len(calls)-index-1)*60
        value = workflow.call(stage, subset, context, count)
        validate_natural(value, subset, originals)
        combined["groups"].extend(value["groups"])
        combined["decisions"].extend(value["decisions"])
    workflow.future_seconds = base_reserve
    # Recheck global naming, membership, boundaries and all original decisions.
    validate_natural(combined, source, original["groups"])
    return combined