"""Helpers for the separate event/time annotation workflow.

The module is deliberately independent of FastAPI and Hugging Face I/O so
the partial-answer rules can be checked locally.
"""

from fastapi import HTTPException
from datetime import datetime, timezone


EVENT_ROLE = "event"
TIME_ROLE = "temporal"
ROLES = (EVENT_ROLE, TIME_ROLE)
QUALIFICATION_RESERVATION_STATUSES = {"in_progress", "retry_ready", "qualified"}
COUNTED_QUALIFICATION_STATUSES = {
    "in_progress", "retry_ready", "qualified",
}


def qualification_passed(state):
    """Only a successfully qualified participant may start the main task."""
    return isinstance(state, dict) and state.get("status") == "qualified"


def qualification_reservation_expired(state, now, ttl_seconds):
    """Whether an unfinished qualification has been idle beyond its allowed time."""
    if (not isinstance(state, dict)
            or state.get("status") not in QUALIFICATION_RESERVATION_STATUSES):
        return False
    timestamp = state.get("updated_at") or state.get("created_at")
    if not isinstance(timestamp, str):
        return False
    try:
        updated_at = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
        if updated_at.tzinfo is None:
            updated_at = updated_at.replace(tzinfo=timezone.utc)
    except ValueError:
        return False
    return (now - updated_at).total_seconds() >= ttl_seconds


def task_assignment_expired(state, now, ttl_seconds):
    """Whether an incomplete split task exceeded its absolute time limit."""
    if not isinstance(state, dict) or state.get("status") != "in_progress":
        return False
    assignment = state.get("assignment")
    timestamp = assignment.get("created_at") if isinstance(assignment, dict) else None
    if not isinstance(timestamp, str):
        timestamp = state.get("created_at")
    if not isinstance(timestamp, str):
        return False
    try:
        started_at = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
        if started_at.tzinfo is None:
            started_at = started_at.replace(tzinfo=timezone.utc)
    except ValueError:
        return False
    return (now - started_at).total_seconds() >= ttl_seconds


def role_occupancy(allocation, qualification_states, now, ttl_seconds):
    """Count task members and live qualification reservations by role."""
    counts = {role: 0 for role in ROLES}
    assigned_ids = set()
    for group in (allocation or {}).get("groups", []):
        members = group.get("members", {})
        if not isinstance(members, dict):
            continue
        for participant_id, role in members.items():
            if role in counts and participant_id not in assigned_ids:
                counts[role] += 1
                assigned_ids.add(participant_id)

    for participant_id, state in (qualification_states or {}).items():
        if (participant_id in assigned_ids
                or not isinstance(state, dict)
                or state.get("task_timeout_released")
                or state.get("status") not in COUNTED_QUALIFICATION_STATUSES
                or qualification_reservation_expired(state, now, ttl_seconds)):
            continue
        role = state.get("annotation_role") if isinstance(state, dict) else None
        if role in counts:
            counts[role] += 1
    return counts


def choose_available_role(counts, capacity_per_role, rng):
    """Balance roles while never assigning beyond a role's remaining capacity."""
    candidates = [role for role in ROLES if counts.get(role, 0) < capacity_per_role]
    if not candidates:
        return None
    least_occupied = min(counts.get(role, 0) for role in candidates)
    return rng.choice([role for role in candidates if counts.get(role, 0) == least_occupied])


def validate_role_responses(item, responses, role):
    """Validate only the fields assigned to this role."""
    if role not in ROLES:
        raise HTTPException(status_code=422, detail="Invalid annotation role")
    if not isinstance(responses, dict):
        raise HTTPException(status_code=422, detail="Invalid annotation responses")

    events = responses.get("events")
    video = responses.get("video")
    if not isinstance(events, dict):
        raise HTTPException(status_code=422, detail="Invalid event responses")
    missing = []
    valid = {"PASS", "FAIL"}

    for event in item.get("events", []):
        event_id = event.get("id")
        answer = events.get(event_id)
        if not isinstance(answer, dict):
            missing.append(event_id)
            continue

        if role == EVENT_ROLE:
            axes = ("v",) if event.get("modality") == "silent_action" else ("a", "v", "avb")
            for axis in axes:
                value = answer.get(axis)
                if not isinstance(value, dict) or value.get("label") not in valid:
                    missing.append(f"{event_id}.{axis}")
        else:
            value = answer.get("tcCount")
            count = value.get("value") if isinstance(value, dict) else None
            try:
                number = int(count)
                valid_count = not isinstance(count, bool) and number >= 0
                if isinstance(count, float):
                    valid_count = valid_count and count.is_integer()
                elif isinstance(count, str):
                    valid_count = valid_count and count.isdigit()
            except (TypeError, ValueError):
                valid_count = False
            if not valid_count:
                missing.append(f"{event_id}.tcCount")

    if role == TIME_ROLE:
        relation = video.get("tcRel") if isinstance(video, dict) else None
        if not isinstance(relation, dict) or relation.get("label") not in valid:
            missing.append("video.tcRel")

    if missing:
        raise HTTPException(
            status_code=422,
            detail=f"Complete all required fields before saving: {', '.join(missing)}",
        )


def required_event_counts(item):
    events = item.get("events", [])
    counts = {event["id"]: 1 for event in events}
    constraint = item.get("temporal_constraint") or {}
    value = constraint.get("value") or {}
    relation = constraint.get("relation_type")
    if relation == "sync_count":
        for event_id in value.get("events", []):
            counts[event_id] = value.get("count", 1)
    elif relation == "order_count":
        counts.update(value.get("event_counts") or {})
    elif relation == "count" and value.get("event"):
        counts[value["event"]] = value.get("count", 1)
    return counts


def aggregate_role_axes(item, responses, role):
    """Return the hidden-GT axes measured by a role for qualification."""
    event_answers = responses["events"]

    def aggregate(axis, include):
        values = [event_answers[event["id"]][axis]["label"]
                  for event in item.get("events", []) if include(event)]
        return "PASS" if values and all(value == "PASS" for value in values) else "FAIL"

    if role == EVENT_ROLE:
        has_audio = lambda event: event.get("modality") != "silent_action"
        return {
            "A": aggregate("a", has_audio),
            "V": aggregate("v", lambda _event: True),
            "AVB": aggregate("avb", has_audio),
        }

    counts = required_event_counts(item)
    matches = all(
        int(event_answers[event_id]["tcCount"]["value"]) == int(expected)
        for event_id, expected in counts.items()
    )
    return {
        "TCREL": responses["video"]["tcRel"]["label"],
        "TCCNT": "PASS" if matches else "FAIL",
    }


def score_role_answer(item, responses, expected, role):
    actual = aggregate_role_axes(item, responses, role)
    correct = sum(actual[axis] == expected[axis] for axis in actual)
    return correct, len(actual)


def merge_role_responses(item, event_responses, temporal_responses):
    """Combine one event part and one temporal part, applying hierarchy rules.

    The input temporal part contains observed raw counts. The returned answer
    follows the existing full-answer shape for downstream analysis.
    """
    events = {}
    event_failed = False
    for event in item.get("events", []):
        event_id = event["id"]
        event_part = event_responses["events"][event_id]
        time_part = temporal_responses["events"][event_id]
        combined = {key: value for key, value in event_part.items() if key in {"a", "v", "avb"}}
        count = time_part["tcCount"]
        failed = any(
            isinstance(value, dict) and value.get("label") == "FAIL"
            for key, value in combined.items() if key in {"a", "v", "avb"}
        )
        if failed:
            event_failed = True
            count = {**count, "value": 0}
        combined["tcCount"] = count
        events[event_id] = combined

    tc_rel = temporal_responses["video"]["tcRel"]
    if event_failed:
        tc_rel = {**tc_rel, "label": "FAIL"}
    return {"events": events, "video": {"tcRel": tc_rel}}