"""Require one model assignment per source alias, then rebuild public partitions."""
from __future__ import annotations

import copy

from suite_contract import Problem, digest

WIRE_INSTRUCTIONS = (
    "OUTPUT REPRESENTATION: Define your natural families in the groups array. "
    "Do not list members there. The assignments object has one REQUIRED key for "
    "EVERY exact input alias. Its value identifies the ZERO-BASED groups array index. "
    "In review passes each assignment instead has family (that index) and variation "
    "(a nonnegative integer naming a related variation set within that family). "
    "Identical content should use the same family and variation. Every group must "
    "be used. Do not omit or rename aliases. This changes only serialization, not "
    "the natural-family objective or the independent assessment. For large-family "
    "decisions use the required review_keys from review_material: those keys identify "
    "exact original memberships, not family names. Never invent decision keys."
)


def decision_keys(context):
    """Index original oversized memberships independently of names or proposal order."""
    original = context.get("original_partition", context.get("natural_partition", {}))
    return {"F-" + digest(sorted(g["members"]))[:20]: sorted(g["members"])
            for g in original.get("groups", []) if len(g["members"]) > 5}


def wire_schema(schema, request, context):
    """Replace freely enumerated memberships with required, exactly keyed assignments."""
    wire = copy.deepcopy(schema)
    group = wire["properties"]["groups"]["items"]
    rich = "variation_sets" in group["properties"]
    for field in ("members", "variation_sets"):
        group["properties"].pop(field, None)
        if field in group["required"]:
            group["required"].remove(field)
    aliases = sorted(c["alias"] for c in request["cases"])
    index = {"type": "integer", "minimum": 0, "maximum": len(aliases) - 1}
    assignment = {"type": "object", "additionalProperties": False,
                  "required": ["family", "variation"],
                  "properties": {"family": index, "variation": index}} if rich else index
    wire["required"].append("assignments")
    wire["properties"]["assignments"] = {
        "type": "object", "additionalProperties": False, "required": aliases,
        "properties": {alias: assignment for alias in aliases}}
    if "decisions" in wire["properties"]:
        item = wire["properties"]["decisions"]["items"]
        del item["properties"]["source_members"]
        item["required"].remove("source_members")
        keys = decision_keys(context)
        wire["properties"]["decisions"] = {
            "type": "object", "additionalProperties": False, "required": sorted(keys),
            "properties": {key: item for key in sorted(keys)}}
    return wire


def decode(value, call, request):
    """Reconstruct groups only after strict alias keys and valid indexes are checked."""
    schema = call["schema"]
    if not isinstance(value, dict) or set(value) != set(schema["required"]):
        raise Problem("invalid_json_result", 502, failure_stage="assignment_envelope")
    aliases = {c["alias"] for c in request["cases"]}
    assignments, groups = value["assignments"], copy.deepcopy(value["groups"])
    if not isinstance(assignments, dict):
        raise Problem("invalid_case_assignments", 502, failure_stage="assignment_keys")
    if set(assignments) != aliases:
        raise Problem("invalid_case_assignments", 502, failure_stage="assignment_keys",
                      missing_case_count=len(aliases - set(assignments)),
                      unknown_case_count=len(set(assignments) - aliases))
    if not isinstance(groups, list) or not 1 <= len(groups) <= len(aliases):
        raise Problem("invalid_case_assignments", 502, failure_stage="assignment_groups")
    required = set(schema["properties"]["groups"]["items"]["required"])
    rich = "common_work" in required
    for group in groups:
        if not isinstance(group, dict) or set(group) != required:
            raise Problem("invalid_json_result", 502, failure_stage="assignment_group_fields")
        group["members"] = []
    variations = [{} for _ in groups]
    for alias in sorted(aliases):
        assignment = assignments[alias]
        if rich:
            if not isinstance(assignment, dict) or set(assignment) != {"family", "variation"}:
                raise Problem("invalid_case_assignments", 502, failure_stage="assignment_value")
            family, variation = assignment["family"], assignment["variation"]
            if type(variation) is not int or not 0 <= variation < len(aliases):
                raise Problem("invalid_variation_assignments", 502)
        else:
            family, variation = assignment, 0
        if type(family) is not int or not 0 <= family < len(groups):
            raise Problem("invalid_case_assignments", 502, failure_stage="assignment_index")
        groups[family]["members"].append(alias)
        variations[family].setdefault(variation, []).append(alias)
    if any(not g["members"] for g in groups):
        raise Problem("invalid_case_assignments", 502, failure_stage="empty_assigned_group")
    if rich:
        for group, sets in zip(groups, variations):
            group["variation_sets"] = [sets[index] for index in sorted(sets)]
    result = {"groups": groups}
    if "decisions" in value:
        decisions, keys = value["decisions"], call["review_keys"]
        if not isinstance(decisions, dict) or set(decisions) != set(keys):
            raise Problem("incomplete_large_family_review", 502, failure_stage="decision_keys")
        if any(not isinstance(item, dict) or set(item) != {"decision", "rationale", "evidence"}
               for item in decisions.values()):
            raise Problem("invalid_review_output", 502)
        result["decisions"] = [{**decisions[key], "source_members": keys[key]} for key in sorted(keys)]
    return result