"""SP28: claim_acks persistence + batch envelope index (D10).

Five methods on top of the new ``ClaimAck`` ORM table:

* ``add_claim_ack`` — insert one link row (manual or auto) and
  publish ``claim_ack_written`` on the bus.
* ``list_acks_for_claim`` — every link row for one claim (per-claim
  only; TA1 batch-level rows are filtered out by the API layer).
* ``list_claims_for_ack`` — every link row for one ack.
* ``find_ack_orphans`` — acks with no resolvable claim (Inbox
  ack-orphans lane).
* ``remove_claim_ack`` — unlink. Publishes ``claim_ack_dropped``.
* ``batch_envelope_index`` — D10 in-memory map of
  ``{Batch.raw_result_json.envelope.control_number (ISA13) OR
  Batch.transaction_set_control_number (ST02)} → batch.id``
  (SP37: populated from two columns; cheap to rebuild, re-built
  once per ingest).

Each mutating method opens its own ``db.SessionLocal()()`` session
so callers don't have to manage session lifecycles. Publishes are
best-effort (wrapped in :func:`_safe_publish`) so a failing bus
subscriber cannot roll back the persisted row.
"""

from __future__ import annotations

import json
import logging
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Iterator

from cyclone import db
from cyclone.db import (
    Ack,
    Batch,
    Claim,
    ClaimAck,
    Ta1Ack,
    Two77caAck,
)

from . import utcnow
from .ui import to_ui_claim_ack
from .write import _sync_publish

if TYPE_CHECKING:
    from cyclone.pubsub import EventBus


log = logging.getLogger(__name__)


def _safe_publish(event_bus: "EventBus | None", kind: str, payload: dict) -> None:
    """Best-effort publish wrapped in try/except.

    Mirrors ``cyclone.store.acks._safe_publish`` — never raises,
    never rolls back the persisted row.
    """
    if event_bus is None:
        return
    try:
        _sync_publish(event_bus, kind, payload)
    except Exception:  # noqa: BLE001
        log.exception("claim_acks store: %s publish failed", kind)


# ---------------------------------------------------------------------------
# D10 batch envelope index
# ---------------------------------------------------------------------------


def batch_envelope_index() -> dict[str, str]:
    """Build a ``{key: batch.id}`` map populated from two columns.

    D10 primary join key (spec §D10). Each 837p batch row contributes
    up to two entries:

    * ``Batch.raw_result_json.envelope.control_number`` (ISA13)
    * ``Batch.transaction_set_control_number`` (ST02, new in SP37)

    Either key resolves to the same ``batch.id``; callers do a single
    ``idx.get(set_control_number)`` lookup. Pass 2 (PCN) is unchanged.

    Built once per ingest from every Batch row whose kind is "837p";
    cost is O(N batches), currently small, so trivial. 835 batches
    are excluded — the column is an 837P-specific join key (see
    ``Batch.transaction_set_control_number`` docstring).
    """
    idx: dict[str, str] = {}
    with db.SessionLocal()() as s:
        rows = (
            s.query(
                Batch.id,
                Batch.raw_result_json,
                Batch.transaction_set_control_number,
            )
            .filter(Batch.kind == "837p")
            .all()
        )
        for bid, raw, stcn in rows:
            env = (raw or {}).get("envelope") or {}
            isa_cn = env.get("control_number")
            if isinstance(isa_cn, str) and isa_cn:
                # First write wins (setdefault) — if ISA13 and ST02
                # ever collide they map to the same batch anyway.
                idx.setdefault(isa_cn, bid)
            if isinstance(stcn, str) and stcn:
                idx.setdefault(stcn, bid)
    return idx


# ---------------------------------------------------------------------------
# Mutators
# ---------------------------------------------------------------------------


def add_claim_ack(
    *,
    claim_id: str | None,
    batch_id: str | None,
    ack_id: int,
    ack_kind: str,
    ak2_index: int | None = None,
    set_control_number: str | None = None,
    set_accept_reject_code: str | None = None,
    linked_by: str = "auto",
    event_bus: "EventBus | None" = None,
    now: datetime | None = None,
) -> ClaimAck:
    """Persist one link row and return it.

    The DB-level 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`` enforces
    idempotency for the per-AK2 case. The 277CA (ak2_index NULL)
    and TA1 (claim_id NULL) paths are deduplicated by application
    code — see :func:`cyclone.claim_acks._link_exists`.

    When ``event_bus`` is provided, publishes one ``claim_ack_written``
    event (with the full :func:`cyclone.store.ui.to_ui_claim_ack`
    payload) so the live-tail subscribers on both
    ``/api/claims/{id}/acks/stream`` and
    ``/api/acks/{kind}/{id}/claims/stream`` see new rows the moment
    they land.
    """
    if ack_kind not in ("999", "277ca", "ta1"):
        raise ValueError(f"add_claim_ack: unknown ack_kind={ack_kind!r}")
    if linked_by not in ("auto", "manual"):
        raise ValueError(f"add_claim_ack: linked_by={linked_by!r}")
    if claim_id is None and batch_id is None:
        raise ValueError(
            "add_claim_ack: at least one of claim_id/batch_id must be set"
        )
    ts = now or utcnow()
    with db.SessionLocal()() as s:
        row = ClaimAck(
            claim_id=claim_id,
            batch_id=batch_id,
            ack_id=ack_id,
            ack_kind=ack_kind,
            ak2_index=ak2_index,
            set_control_number=set_control_number,
            set_accept_reject_code=set_accept_reject_code,
            linked_at=ts,
            linked_by=linked_by,
        )
        s.add(row)
        s.commit()
        s.refresh(row)
        row_id = row.id

    if event_bus is not None:
        with db.SessionLocal()() as s:
            payload = to_ui_claim_ack(s.get(ClaimAck, row_id))
        _safe_publish(event_bus, "claim_ack_written", payload)

    return row


def remove_claim_ack(
    link_id: int,
    *,
    event_bus: "EventBus | None" = None,
) -> bool:
    """Unlink. Returns ``True`` when a row was deleted.

    Publishes ``claim_ack_dropped`` with ``{"id", "claim_id"}`` so
    the live-tail subscribers can remove the link from their local
    store. Does NOT touch ``Claim.state`` — the unlink is purely
    about the link row, per spec §D6.
    """
    with db.SessionLocal()() as s:
        row = s.get(ClaimAck, link_id)
        if row is None:
            return False
        payload = {
            "id": row.id,
            "claim_id": row.claim_id,
            "ack_id": row.ack_id,
            "ack_kind": row.ack_kind,
        }
        s.delete(row)
        s.commit()

    if event_bus is not None:
        _safe_publish(event_bus, "claim_ack_dropped", payload)
    return True


# ---------------------------------------------------------------------------
# Readers
# ---------------------------------------------------------------------------


def list_acks_for_claim(claim_id: str) -> list[ClaimAck]:
    """Return every ClaimAck row for one claim (newest first)."""
    with db.SessionLocal()() as s:
        return (
            s.query(ClaimAck)
            .filter(ClaimAck.claim_id == claim_id)
            .order_by(ClaimAck.id.desc())
            .all()
        )


def list_claims_for_ack(kind: str, ack_id: int) -> list[ClaimAck]:
    """Return every ClaimAck row for one ack (any kind).

    For 999 / 277CA: returns 0..N rows (one per AK2 / ClaimStatus).
    For TA1: returns 0..1 row (envelope-level, populated batch_id).
    """
    if kind not in ("999", "277ca", "ta1"):
        raise ValueError(f"list_claims_for_ack: unknown kind={kind!r}")
    with db.SessionLocal()() as s:
        return (
            s.query(ClaimAck)
            .filter(
                ClaimAck.ack_kind == kind,
                ClaimAck.ack_id == ack_id,
            )
            .order_by(ClaimAck.id.asc())
            .all()
        )


# ---------------------------------------------------------------------------
# Orphan detection (Inbox "Ack orphans" lane — spec §D7)
# ---------------------------------------------------------------------------


def find_ack_orphans(kind: str) -> list[dict]:
    """Return acks with no resolvable Claim row of their own kind.

    Used by the Inbox ack-orphans lane for the operator's manual
    reconciliation flow. The downstream ``/api/inbox/ack-orphans``
    endpoint calls this and returns the rendered shape.

    "Orphan" means: for 999 / 277CA, ``ack_kind=kind AND ack_id IN
    (acks_with_no_claim_acks_link)``. For TA1, "orphan" means the
    TA1 row exists but no Batch with matching sender/receiver was
    resolved (so no link row was created).

    Output dict shape (one row per orphan ack, rendered by
    :func:`to_ui_claim_ack`-style serialization):

      * ``kind``       — "999" / "277ca" / "ta1"
      * ``ack_id``     — the ack row's id
      * ``control_number`` — for 999/277CA, envelope.control_number;
                             for TA1, ta1.control_number
      * ``set_control_numbers`` — empty list when no claims match
      * ``raw_summary`` — flat copy of the ack's UI shape for the
                          lane-header counts
    """
    if kind not in ("999", "277ca", "ta1"):
        raise ValueError(f"find_ack_orphans: unknown kind={kind!r}")

    out: list[dict] = []
    if kind == "999":
        ack_table = Ack
    elif kind == "277ca":
        ack_table = Two77caAck
    else:
        ack_table = Ta1Ack

    with db.SessionLocal()() as s:
        # Every ack row of the given kind; orphan when NO claim_acks
        # link was created for it. The auto-linker emits one claim_acks
        # row per AK2 even when the AK2 is rejected, so a 999 with at
        # least one linked AK2 is never an orphan.
        for ack_row in s.query(ack_table).order_by(ack_table.id.desc()).all():
            count = (
                s.query(ClaimAck)
                .filter(ClaimAck.ack_kind == kind,
                        ClaimAck.ack_id == ack_row.id)
                .count()
            )
            if count > 0:
                continue
            out.append({
                "kind": kind,
                "ack_id": ack_row.id,
                "control_number": _ack_control_number(ack_row, kind),
                "parsed_at": (
                    ack_row.parsed_at.isoformat().replace(
                        "+00:00", "Z"
                    ) if ack_row.parsed_at else None
                ),
            })
    return out


def _iter_orphan_999_st02s() -> Iterator[tuple[str, int]]:
    """Yield ``(set_control_number, orphan_count)`` for each distinct orphan 999.

    An orphan 999 has zero ``claim_acks`` rows tied to its id (per
    :func:`find_ack_orphans`). The ST02 is the source 837's
    ``transaction_set_control_number``, extracted from the 999's
    ``set_responses[0].set_control_number`` field in ``raw_json``.

    Used by :meth:`CycloneStore.find_ack_orphan_st02_summary` and
    :meth:`CycloneStore.reconcile_orphan_st02s` (sp38) to enumerate
    orphan ST02s without re-implementing the orphan-detection query.

    Skips 999s whose ``raw_json`` is malformed or has no
    ``set_responses`` entry — those are unprocessable regardless of
    whether they have a matching batch row.

    999-only. 277ca / ta1 orphans are tracked separately by the
    store and are not aggregated here (their "ST02" semantics differ
    from 999: it's the interchange control number, not the source
    837's ST02).
    """
    with db.SessionLocal()() as s:
        # LEFT OUTER JOIN keeps all acks; orphan when no claim_acks row
        # exists for (ack_kind='999', ack_id=acks.id).
        rows = (
            s.query(Ack)
            .outerjoin(
                ClaimAck,
                (ClaimAck.ack_kind == "999") & (ClaimAck.ack_id == Ack.id),
            )
            .filter(ClaimAck.id.is_(None))
            .all()
        )
        counts: dict[str, int] = {}
        for ack_row in rows:
            st02 = _extract_999_st02(ack_row)
            if st02 is None:
                continue
            counts[st02] = counts.get(st02, 0) + 1
    # Unsorted on purpose — the caller re-sorts by (-ack_count, st02)
    # so the heaviest backlog surfaces first in CLI output. Sorting
    # here would just be wasted CPU (and risk a different order if
    # the caller's key changes).
    yield from counts.items()


def _extract_999_st02(ack_row: Ack) -> str | None:
    """Return the source 837's ST02 from a 999 ack row, or ``None``.

    Reads ``raw_json`` and pulls ``set_responses[0].set_control_number``.
    Returns ``None`` if the JSON is malformed, ``set_responses`` is
    empty, or the field is missing — callers should skip these rows
    rather than treating them as orphans (they're unprocessable, not
    just orphaned).

    Note: the ``raw_json`` column is a SQLAlchemy JSON type so the
    ORM hands it back as a Python dict, not a string. We accept
    either form (dict or str) so this helper is robust to direct
    SQLAlchemy access and to raw sqlite3 row access.
    """
    raw = ack_row.raw_json
    if not raw:
        return None
    if isinstance(raw, dict):
        parsed = raw
    elif isinstance(raw, (str, bytes, bytearray)):
        try:
            parsed = json.loads(raw)
        except (TypeError, ValueError):
            return None
    else:
        return None
    srs = parsed.get("set_responses") or []
    if not srs:
        return None
    st02 = srs[0].get("set_control_number")
    return str(st02) if st02 else None


def _ack_control_number(ack_row, kind: str) -> str:
    """Best-effort control-number lookup for an ack row of any kind.

    Per-kind sources (preserved across the SP38 refactor of
    :func:`find_ack_orphans`):

      * ``999`` — 999 ORM row has no dedicated control_number column;
        we re-derive it from ``raw_json.envelope.control_number``
        (the same source :func:`cyclone.store.ui.to_ui_ack` uses).
        Accepts both dict (post-ORM hydration) and str (raw sqlite3
        row or freshly inserted JSON string).
      * ``277ca`` / ``ta1`` — both carry the control number in a
        dedicated ORM column (``ack_row.control_number``). Reading
        from ``raw_json`` would yield empty strings.

    Returns ``""`` (empty string) when the source field is missing
    so the JSON renderer still emits a stable shape.
    """
    if kind == "999":
        raw = ack_row.raw_json
        if raw is None or raw == "":
            return ""
        if isinstance(raw, dict):
            env = raw.get("envelope") or {}
            return env.get("control_number") or ""
        if isinstance(raw, (str, bytes, bytearray)):
            try:
                parsed = json.loads(raw)
            except (TypeError, ValueError):
                return ""
            env = parsed.get("envelope") or {}
            return env.get("control_number") or ""
        return ""
    # 277ca / ta1 — control_number is an ORM column on both tables.
    return getattr(ack_row, "control_number", "") or ""


__all__ = [
    "add_claim_ack",
    "batch_envelope_index",
    "find_ack_orphans",
    "list_acks_for_claim",
    "list_claims_for_ack",
    "remove_claim_ack",
]
