"""Derived continuity grouping for immutable route-association evidence."""

from __future__ import annotations

import math

from mountain_twin.trail_character.contracts import (
    AssociationContinuityKind,
    AssociationContinuitySpan,
    AssociationState,
    AssociationStateDistance,
    RouteAssociationContinuityPolicy,
    RouteAssociationContinuityResult,
    RouteAssociationResult,
)

DEFAULT_ASSOCIATION_CONTINUITY_POLICY = RouteAssociationContinuityPolicy(
    policy_id="mt_route_association_continuity_rnd",
    version="v0_1",
    threshold_origin="MT_RND_CONTINUITY_POLICY",
    maximum_raw_unmatched_gap_m=5.0,
)


def derive_association_continuity(
    association: RouteAssociationResult,
    *,
    policy: RouteAssociationContinuityPolicy = DEFAULT_ASSOCIATION_CONTINUITY_POLICY,
) -> RouteAssociationContinuityResult:
    """Group only identical MATCHED evidence across short raw UNMATCHED gaps.

    The raw association remains attached unchanged. A derived span containing a
    gap is not a MATCHED state: its ``raw_state_distances`` exposes the exact
    MATCHED and UNMATCHED distances. AMBIGUOUS, UNKNOWN, candidate changes,
    and longer UNMATCHED gaps always remain separate spans.
    """
    if not _valid_policy(policy):
        raise ValueError("association continuity policy is invalid")

    raw_spans = association.spans
    continuity_spans: list[AssociationContinuitySpan] = []
    index = 0
    while index < len(raw_spans):
        end_index = _extend_same_candidate_group(raw_spans, index, policy)
        continuity_spans.append(_continuity_span(raw_spans, index, end_index))
        index = end_index + 1
    return RouteAssociationContinuityResult(association, policy, tuple(continuity_spans))


def _extend_same_candidate_group(
    spans, start_index: int, policy: RouteAssociationContinuityPolicy
) -> int:
    start = spans[start_index]
    if start.state is not AssociationState.MATCHED:
        return start_index

    end_index = start_index
    while end_index + 2 < len(spans):
        gap = spans[end_index + 1]
        next_span = spans[end_index + 2]
        if (
            gap.state is not AssociationState.UNMATCHED
            or _span_length_m(gap) > policy.maximum_raw_unmatched_gap_m
            or next_span.state is not AssociationState.MATCHED
            or next_span.candidate_feature_ids != start.candidate_feature_ids
        ):
            break
        end_index += 2
    return end_index


def _continuity_span(raw_spans, start_index: int, end_index: int) -> AssociationContinuitySpan:
    included = raw_spans[start_index : end_index + 1]
    has_short_unmatched_gap = len(included) > 1
    first, last = included[0], included[-1]
    state_distances = tuple(
        AssociationStateDistance(
            state, sum(_span_length_m(item) for item in included if item.state is state)
        )
        for state in AssociationState
        if any(item.state is state for item in included)
    )
    return AssociationContinuitySpan(
        kind=(
            AssociationContinuityKind.SAME_CANDIDATE_WITH_SHORT_RAW_UNMATCHED_GAPS
            if has_short_unmatched_gap
            else AssociationContinuityKind.EXACT_RAW_ASSOCIATION
        ),
        start_route_distance_m=first.start_route_distance_m,
        end_route_distance_m=last.end_route_distance_m,
        candidate_feature_ids=first.candidate_feature_ids,
        raw_span_start_index=start_index,
        raw_span_end_index=end_index,
        raw_state_distances=state_distances,
        reason_codes=(
            ("SAME_CANDIDATE_SHORT_RAW_UNMATCHED_GAPS_RETAINED",)
            if has_short_unmatched_gap
            else ("RAW_ASSOCIATION_SPAN_UNCHANGED",)
        ),
    )


def _span_length_m(span) -> float:
    return span.end_route_distance_m - span.start_route_distance_m


def _valid_policy(policy: RouteAssociationContinuityPolicy) -> bool:
    return (
        policy.threshold_origin == "MT_RND_CONTINUITY_POLICY"
        and math.isfinite(policy.maximum_raw_unmatched_gap_m)
        and policy.maximum_raw_unmatched_gap_m >= 0
    )