"""999 ACK → claim state transitions for the Inbox Rejected lane.

SP6 T3: apply_999_rejections.
For each set response with code R or E, look up the matching claim and
move it to REJECTED. Idempotent on already-rejected claims.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Callable

from sqlalchemy.orm import Session

from cyclone.db import Claim, ClaimState


@dataclass
class Apply999Result:
    matched: list[str] = field(default_factory=list)   # claim ids transitioned
    orphans: list[str] = field(default_factory=list)   # PCNs we couldn't resolve


def _build_reason(code: str, segment_error_count: int) -> str:
    parts = [f"999 AK5 set-level {code}"]
    if segment_error_count:
        parts.append(f"{segment_error_count} segment error(s)")
    return "; ".join(parts)


def apply_999_rejections(
    session: Session,
    parsed_999,
    *,
    claim_lookup: Callable[[str], Claim | None],
    batch_envelope_index: dict[str, list[str]] | None = None,
) -> Apply999Result:
    """For each set response with code R, E, or X, look up the matching claim and
    move it to REJECTED. Idempotent on already-rejected claims.

    Args:
        session: SQLAlchemy session.
        parsed_999: a ParseResult999 (or any object with .set_responses).
        claim_lookup: callable from patient_control_number → Claim or None.
            Legacy fallback; rarely hits when batch_envelope_index is present.
        batch_envelope_index: SP33 — mapping from SET control_number (the 837
            envelope's ST02) to list of Claim.id for the claims in that SET.
            Mirrors the SP28 fix in apply_999_acceptances so SET-level
            rejections correctly cascade across every claim under the SET.

    Returns:
        Apply999Result with lists of matched claim ids and orphan PCNs.
    """
    result = Apply999Result()
    now = datetime.now(timezone.utc)
    index = batch_envelope_index or {}

    for sr in parsed_999.set_responses:
        code = sr.set_accept_reject.code
        if code not in ("R", "E", "X"):
            continue

        # SP33: prefer batch_envelope_index (SCN -> [claim_id]) so a SET-level
        # rejection correctly flips every claim in the SET. Fall back to
        # the legacy claim_lookup when the index is empty for this SCN.
        candidate_ids = index.get(sr.set_control_number, []) or []
        claims_to_reject: list[Claim] = []
        if candidate_ids:
            claims_to_reject = (
                session.query(Claim)
                .filter(Claim.id.in_(candidate_ids))
                .all()
            )
        else:
            legacy = claim_lookup(sr.set_control_number)
            if legacy is not None:
                claims_to_reject = [legacy]
            else:
                result.orphans.append(sr.set_control_number)
                continue

        for claim in claims_to_reject:
            if claim.state == ClaimState.REJECTED:
                # Idempotent: don't double-mutate.
                continue
            claim.state = ClaimState.REJECTED
            claim.state_changed_at = now
            claim.rejected_at = now
            claim.rejection_reason = _build_reason(
                code, len(sr.segment_errors or [])
            )
            result.matched.append(claim.id)

    if result.matched or result.orphans:
        session.commit()

    return result
