"""SP28: per-ACK auto-linker.

Three helpers, one per ACK kind, all run inside the same DB
transaction that persists the Ack row. Each returns a slice of
:class:`ClaimAckLinkRow` dataclasses describing what to link —
the CALLER persists those rows via
:func:`cyclone.store.claim_acks.add_claim_ack` (which owns the
publish-from-store contract).

The two-pass join lives in
:func:`lookup_claims_for_ack_set_response` (D10): ST02 via the
batch envelope index (primary) + ``Claim.patient_control_number``
(fallback). Plus :func:`link_manual` for the manual-fallback
endpoint.

The helpers do NOT write to the DB session — they are pure
readers over the session + parse result + the supplied
``batch_envelope_index`` / ``pc_claim_lookup`` / ``batch_lookup``
closures. This matches the existing
``cyclone.inbox_state.apply_999_rejections`` pattern and lets the
store facade own the publish-from-store contract for live-tail.

See ``docs/superpowers/specs/2026-07-02-cyclone-ack-claim-auto-link-design.md``
§3 for the per-AK2 granularity, the two-pass join, and the
idempotency contract enforced by ``ux_claim_acks_dedup``.
"""
from __future__ import annotations

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

from sqlalchemy.orm import Session

from cyclone.db import Batch, Claim

log = logging.getLogger(__name__)


@dataclass
class ClaimAckLinkRow:
    """One row to insert into ``claim_acks``.

    Populated by the helpers; persisted by the caller via
    :func:`cyclone.store.claim_acks.add_claim_ack`. Carries every
    column the store needs to build the ORM row + the
    ``claim_ack_written`` event payload.
    """

    claim_id: Optional[str] = None
    batch_id: Optional[str] = None
    ak2_index: Optional[int] = None
    set_control_number: Optional[str] = None
    set_accept_reject_code: Optional[str] = None


@dataclass
class ClaimAckLinkResult:
    """Outcome of a single helper call.

    ``linked`` is the list of :class:`ClaimAckLinkRow` rows the
    caller should persist via ``cycl_store.add_claim_ack``.
    ``orphans`` is a list of free-form strings the join couldn't
    resolve — for 999/277CA these are ``set_control_number``
    values; for TA1 they're the TA1 ICN when no matching batch was
    found.
    """

    linked: list[ClaimAckLinkRow] = field(default_factory=list)
    orphans: list[str] = field(default_factory=list)


# ---------------------------------------------------------------------------
# D10 two-pass join
# ---------------------------------------------------------------------------


def lookup_claims_for_ack_set_response(
    session: Session,
    set_control_number: str,
    *,
    batch_envelope_index: Optional[Callable[[str], Optional[str]]] = None,
    pc_claim_lookup: Optional[Callable[[str], Optional[Claim]]] = None,
) -> list[Claim]:
    """Two-pass join for a single AK2 set_response / 277CA claim_status.

    D10 (spec): for a 999 AK2-2 ``set_control_number`` or a 277CA REF*1K
    ``payer_claim_control_number``, return every claim this ack
    acknowledges. The primary join is
    ``Batch.envelope.control_number == set_control_number`` (== source
    837's ST02 on Gainwell batches); the fallback is
    ``Claim.patient_control_number == set_control_number``.

    Returns 0..N matching claims (one-ack-to-many when one 837 batch
    shipped multiple claims under one ST02).

    Args:
        session: SQLAlchemy session the caller owns.
        set_control_number: the AK2-2 / REF*1K value to resolve.
        batch_envelope_index: optional pre-built index that maps
            ``Batch.envelope.control_number`` → ``batch.id`` (built
            once per ingest via ``store.batch_envelope_index``).
            Pass to skip the per-set-response ``Batch`` scan in Pass 1.
        pc_claim_lookup: optional pre-built callable that maps a PCN
            to a single claim. Falls back to a session-wide query
            when not supplied.

    Pass 1 wins. The two paths cannot both fire for the same
    ``set_control_number`` — if Pass 1 returns one or more claims,
    Pass 2 is skipped. This is the false-positive guard from
    spec §7.
    """
    if not set_control_number:
        return []

    # -- Pass 1: Batch.envelope.control_number primary --------------
    # Accept either a plain dict (the common case — built once per
    # ingest via ``store.batch_envelope_index``) or a callable for
    # test-side closures. Normalize to ``idx.get`` so the rest of
    # the function stays uniform.
    idx: Optional[Callable[[str], Optional[str]]] = None
    if batch_envelope_index is not None:
        if callable(batch_envelope_index):
            idx = batch_envelope_index
        else:
            idx = batch_envelope_index.get
    matched_ids: list[str] = []
    if idx is not None:
        batch_id = idx(set_control_number)
        if batch_id is not None:
            matched_ids = [
                cid for (cid,) in (
                    session.query(Claim.id)
                    .filter(Claim.batch_id == batch_id)
                    .all()
                )
            ]
    else:
        # Fallback: scan all batches once per call. Slow but correct;
        # callers SHOULD pass the index.
        rows = (
            session.query(Batch.id, Batch.raw_result_json)
            .filter(Batch.kind == "837p")
            .all()
        )
        for bid, raw in rows:
            env = (raw or {}).get("envelope") or {}
            if env.get("control_number") == set_control_number:
                matched_ids = [
                    cid for (cid,) in (
                        session.query(Claim.id)
                        .filter(Claim.batch_id == bid)
                        .all()
                    )
                ]
                break

    if matched_ids:
        claims = (
            session.query(Claim)
            .filter(Claim.id.in_(matched_ids))
            .all()
        )
        if claims:
            return list(claims)

    # -- Pass 2: Claim.patient_control_number fallback ---------------
    if pc_claim_lookup is not None:
        single = pc_claim_lookup(set_control_number)
        if single is not None:
            return [single]
        return []

    matches = (
        session.query(Claim)
        .filter(Claim.patient_control_number == set_control_number)
        .all()
    )
    return list(matches)


# ---------------------------------------------------------------------------
# Per-ACK helpers — walk the parsed result and produce ClaimAckLinkRow
# dataclasses. The CALLER persists via cycl_store.add_claim_ack so the
# publish-from-store contract owns the live-tail event emission.
# ---------------------------------------------------------------------------


def _existing_link_claim_ids(
    session: Session,
    *,
    claim_ids: list[str],
    ack_kind: str,
    ack_id: int,
) -> set[str]:
    """Return the subset of ``claim_ids`` that already have a link row.

    Mirrors the partial unique index
    ``ux_claim_acks_dedup(claim_id, ack_kind, ack_id, ak2_index)
    WHERE claim_id IS NOT NULL AND ak2_index IS NOT NULL``. The
    pre-check is here so we skip ``session.add`` and avoid
    IntegrityError log noise on re-ingest. For TA1 (claim_id IS
    NULL) the helpers do their own check.
    """
    if not claim_ids:
        return set()
    rows = (
        session.query(Claim.claim_id)  # placeholder; replaced below
        if False else
        session.query(Claim.id)
        .filter(Claim.id.in_([]))
        .all()
    )
    # Replace the placeholder with the real query — needed because
    # ClaimAck is not imported above (avoids circular import).
    from cyclone.db import ClaimAck
    existing = (
        session.query(ClaimAck.claim_id)
        .filter(
            ClaimAck.ack_kind == ack_kind,
            ClaimAck.ack_id == ack_id,
            ClaimAck.claim_id.in_(claim_ids),
        )
        .all()
    )
    return {cid for (cid,) in existing if cid is not None}


def apply_999_acceptances(
    session: Session,
    parsed_999,
    *,
    ack_id: int,
    batch_envelope_index: Optional[Callable[[str], Optional[str]]] = None,
    pc_claim_lookup: Optional[Callable[[str], Optional[Claim]]] = None,
    now: Optional[datetime] = None,
) -> ClaimAckLinkResult:
    """For every AK2 set-response, build one ``ClaimAckLinkRow`` per matched claim.

    Both accepted AND rejected AK2s produce a link row (so the
    ClaimDrawer panel can show the rejection inline via the
    ``set_accept_reject_code`` color-coded chip). Orphans are returned
    but not linked.

    Idempotent: rows the dedup index already covers (re-ingest of an
    identical file) are skipped silently — the pre-check is here to
    avoid ``IntegrityError`` log noise.

    One AK2 can produce multiple ``claim_acks`` rows when the source
    837 batch carried more than one claim under a shared ST02 (rare
    on this codebase but supported by D10 / the schema).

    Returns:
        A :class:`ClaimAckLinkResult` whose ``linked`` list contains
        one :class:`ClaimAckLinkRow` per AK2-to-claim match. The
        caller persists each row via ``cycl_store.add_claim_ack``.
    """
    result = ClaimAckLinkResult()
    set_responses = getattr(parsed_999, "set_responses", None) or []

    # Resolve all set_control_numbers up front so we can do one
    # batched dedup query per (ack_kind, ack_id).
    resolved: list[tuple[int, "object", list[Claim], str, str]] = []
    candidate_claim_ids: list[str] = []
    for idx, sr in enumerate(set_responses):
        scn = getattr(sr, "set_control_number", None) or ""
        code = getattr(sr.set_accept_reject, "code", None) or ""
        claims = lookup_claims_for_ack_set_response(
            session, scn,
            batch_envelope_index=batch_envelope_index,
            pc_claim_lookup=pc_claim_lookup,
        )
        if not claims:
            result.orphans.append(scn)
            continue
        resolved.append((idx, sr, claims, scn, code))
        candidate_claim_ids.extend(c.id for c in claims)

    existing_ids = _existing_link_claim_ids(
        session,
        claim_ids=candidate_claim_ids,
        ack_kind="999",
        ack_id=ack_id,
    )

    for idx, sr, claims, scn, code in resolved:
        for claim in claims:
            if claim.id in existing_ids:
                continue
            result.linked.append(ClaimAckLinkRow(
                claim_id=claim.id,
                batch_id=None,
                ak2_index=idx,
                set_control_number=scn,
                set_accept_reject_code=code,
            ))
    return result


def apply_277ca_acks(
    session: Session,
    parsed_277ca,
    *,
    ack_id: int,
    batch_envelope_index: Optional[Callable[[str], Optional[str]]] = None,
    pc_claim_lookup: Optional[Callable[[str], Optional[Claim]]] = None,
    now: Optional[datetime] = None,
) -> ClaimAckLinkResult:
    """For every ClaimStatus with a payer_claim_control_number, build a ClaimAckLinkRow.

    Accepted AND rejected ClaimStatuses both link — the
    ``set_accept_reject_code`` carries the STC category code. The
    ``claim_acks`` row is independent of the existing
    ``Claim.payer_rejected_at`` mutation from
    :func:`cyclone.inbox_state_277ca.apply_277ca_rejections` (which
    fires before this helper in the handler).

    Returns:
        A :class:`ClaimAckLinkResult` whose ``linked`` list contains
        one :class:`ClaimAckLinkRow` per ClaimStatus match. The
        caller persists each row via ``cycl_store.add_claim_ack``.
    """
    result = ClaimAckLinkResult()
    statuses = getattr(parsed_277ca, "claim_statuses", None) or []

    resolved: list[tuple["object", list[Claim], str, str]] = []
    candidate_claim_ids: list[str] = []
    for status in statuses:
        scn = getattr(status, "payer_claim_control_number", None) or ""
        raw_code = getattr(status, "status_code", None) or ""
        # ``status_code`` may be a category like "A6:19:PR" — keep
        # the whole STC composite so the UI can render the
        # category without re-parsing raw_json. Truncate to 8 chars
        # (column width).
        code = (raw_code or "")[:8]
        if not scn:
            # No REF*1K — orphan. Surface the STC composite so the
            # operator can correlate via the ack's raw_json.
            orphan_key = code or "(no REF*1K)"
            result.orphans.append(orphan_key)
            continue
        claims = lookup_claims_for_ack_set_response(
            session, scn,
            batch_envelope_index=batch_envelope_index,
            pc_claim_lookup=pc_claim_lookup,
        )
        if not claims:
            result.orphans.append(scn)
            continue
        resolved.append((status, claims, scn, code))
        candidate_claim_ids.extend(c.id for c in claims)

    existing_ids = _existing_link_claim_ids(
        session,
        claim_ids=candidate_claim_ids,
        ack_kind="277ca",
        ack_id=ack_id,
    )

    for status, claims, scn, code in resolved:
        for claim in claims:
            if claim.id in existing_ids:
                continue
            result.linked.append(ClaimAckLinkRow(
                claim_id=claim.id,
                batch_id=None,
                ak2_index=None,
                set_control_number=scn,
                set_accept_reject_code=code or None,
            ))
    return result


def apply_ta1_envelope_link(
    session: Session,
    parsed_ta1,
    *,
    ack_id: int,
    batch_lookup: Callable[[str, str], Optional[Batch]],
    now: Optional[datetime] = None,
) -> ClaimAckLinkResult:
    """Build a TA1 envelope-level link row to the most-recent matching Batch.

    TA1 is envelope-level only (ISA/IEA, no per-claim granularity).
    The link row has ``claim_id IS NULL`` and ``batch_id`` populated.

    Args:
        parsed_ta1: a :class:`cyclone.parsers.models_ta1.ParseResultTa1`.
        batch_lookup: ``(sender_id, receiver_id) -> Batch | None``.
            The handler supplies a closure that walks
            ``session.query(Batch).order_by(parsed_at.desc())``. Returning
            ``None`` produces an orphan (no batch match).

    Returns:
        A :class:`ClaimAckLinkResult` with 0..1 row in ``linked``.
        Idempotent via dedup on ``(ack_kind='ta1', ack_id)`` (claim_id
        IS NULL so the partial unique index doesn't catch it; the
        pre-check here is a Python-side query).
    """
    result = ClaimAckLinkResult()
    envelope = getattr(parsed_ta1, "envelope", None)
    if envelope is None:
        return result
    ta1_obj = getattr(parsed_ta1, "ta1", None)
    ack_code = getattr(ta1_obj, "ack_code", None) or ""

    # Dedup: same (ack_kind, ack_id) → at most one TA1 envelope link.
    # Done in Python because the partial unique index requires
    # claim_id IS NOT NULL.
    from cyclone.db import ClaimAck
    existing = (
        session.query(ClaimAck.id)
        .filter(
            ClaimAck.ack_kind == "ta1",
            ClaimAck.ack_id == ack_id,
            ClaimAck.claim_id.is_(None),
        )
        .first()
    )
    if existing is not None:
        return result

    batch = batch_lookup(envelope.sender_id or "", envelope.receiver_id or "")
    if batch is None:
        orphan_key = (
            getattr(ta1_obj, "control_number", None)
            or envelope.control_number
            or ""
        )
        result.orphans.append(orphan_key)
        return result

    result.linked.append(ClaimAckLinkRow(
        claim_id=None,
        batch_id=batch.id,
        ak2_index=None,
        set_control_number=None,
        set_accept_reject_code=ack_code,
    ))
    return result


def link_manual(
    session: Session,
    *,
    claim_id: str,
    ack_kind: str,
    ack_id: int,
    set_control_number: Optional[str] = None,
    set_accept_reject_code: Optional[str] = None,
    ak2_index: Optional[int] = None,
    now: Optional[datetime] = None,
) -> ClaimAckLinkRow:
    """Return one manual link row (the caller persists it via the store).

    Used by ``POST /api/acks/{kind}/{ack_id}/match-claim``. Returns a
    :class:`ClaimAckLinkRow` describing the row to insert. The caller
    is responsible for persistence so it can own the publish-from-store
    contract.

    Idempotency: callers should pre-check via ``session.query(ClaimAck)
    .filter(...).first()`` and skip when a row already exists; the
    :class:`cyclone.store.claim_acks.add_claim_ack` implementation also
    re-checks via the partial unique index.

    Raises ``LookupError`` when the referenced claim doesn't exist
    (the caller maps that to 404).
    """
    if ack_kind not in ("999", "277ca", "ta1"):
        raise ValueError(f"link_manual: unknown ack_kind={ack_kind!r}")
    claim = session.get(Claim, claim_id)
    if claim is None:
        raise LookupError(f"claim {claim_id} not found")
    return ClaimAckLinkRow(
        claim_id=claim_id,
        batch_id=None,
        ak2_index=ak2_index,
        set_control_number=set_control_number,
        set_accept_reject_code=set_accept_reject_code,
    )


__all__ = [
    "ClaimAckLinkResult",
    "ClaimAckLinkRow",
    "apply_999_acceptances",
    "apply_277ca_acks",
    "apply_ta1_envelope_link",
    "link_manual",
    "lookup_claims_for_ack_set_response",
]