"""Pydantic v2 models for parsed 999 (Implementation Acknowledgment) files.

Mirrors the X12 005010X231A1 segment shape:

- ``AcknowledgmentHeader`` (AK1 / AK2) — identifies the source set/functional group
- ``SetFunctionalGroupResponse`` (AK2 + AK3*/AK4*/AK5) — one per set
- ``FunctionalGroupAck`` (AK1 + AK9) — one per functional group
- ``ParseResult999`` (batch)

Per spec §4.1, the leaner shape below is sufficient for v1. The richer
``FunctionalGroupResponse`` (with ``SetHeader``/``ElementError`` types)
described in the spec is folded into the flat ``SetFunctionalGroupResponse``
that the parser/serializer actually construct.

We copy the ``_Base`` model from :mod:`cyclone.parsers.models` so the
``date`` -> ``str`` JSON serializer is shared with the 837P / 835 models.
That keeps FastAPI responses consistent across transaction sets.
"""
from __future__ import annotations

from datetime import date
from typing import Literal

from pydantic import BaseModel, ConfigDict, Field, model_serializer

from cyclone.parsers.models import Address, BatchSummary, Envelope, ValidationIssue, ValidationReport


# --------------------------------------------------------------------------- #
# Shared base
# --------------------------------------------------------------------------- #


class _Base(BaseModel):
    """Shared Pydantic base; matches the 837P / 835 models for JSON consistency."""

    model_config = ConfigDict(extra="ignore", str_strip_whitespace=True)

    @model_serializer(mode="wrap")
    def _serialize(self, handler):  # type: ignore[no-untyped-def]
        data = handler(self)
        for key, value in data.items():
            if isinstance(value, date):
                data[key] = value.isoformat()
        return data


# --------------------------------------------------------------------------- #
# 999 header / segment shapes
# --------------------------------------------------------------------------- #


class AcknowledgmentHeader(_Base):
    """AK1 / AK2 — identifies the source functional group or set.

    AK101 = functional_id_code (e.g. "HC" for health-care claim).
    AK102 = group_control_number (from GS06 for AK1, from ST02 for AK2).
    For an AK2, ``functional_id_code`` carries the transaction set ID
    (e.g. "837") rather than the functional group code; both elements
    are required by the spec.
    """

    functional_id_code: str
    group_control_number: str


class SegmentContext(_Base):
    """AK3 — segment context for an error report.

    AK301 = segment_id (e.g. "CLM"), AK302 = segment_position,
    AK303 = loop_id (optional), AK304 = implementation_convention_ref
    (optional). Per the X12 999 spec, AK304 is rarely used in practice
    so we keep it optional.
    """

    segment_id: str
    segment_position: int
    loop_id: str | None = None
    implementation_convention_ref: str | None = None


class SegmentError(_Base):
    """AK3 + AK4 — a single segment- or element-level error report.

    AK4-1 carries the element position (1-based) when the error
    references a specific element. When AK4 is omitted (segment-level
    only), element_position is ``None``. AK4-2 = element_reference and
    AK4-3 = error_code.
    """

    context: SegmentContext
    error_code: str
    element_position: int | None = None
    element_reference: int | None = None


class SetAcceptReject(_Base):
    """AK5 — per-set accept/reject code.

    Standard codes per X12 005010X231A1:
      A = Accepted
      R = Rejected
      E = Accepted with errors
      W = Warning
      X = Rejected (all sets in the functional group)
    """

    code: Literal["A", "R", "E", "W", "X"]


class SetFunctionalGroupResponse(_Base):
    """AK2 + (optional AK3*/AK4*) + AK5 — one per set in the source batch.

    We surface the source set's ST01 (transaction_set_identifier) and
    ST02 (set_control_number) on the response so callers can correlate
    ACKs back to the original 837/835 sets without a second lookup.

    ``ak2`` carries the AK201/AK202 values. The AK2's
    ``functional_id_code`` field holds the transaction set ID (e.g.
    "837"); this matches the spec's reuse of the AK1 element shape.
    """

    ak2: AcknowledgmentHeader
    set_control_number: str
    transaction_set_identifier: str
    segment_errors: list[SegmentError] = Field(default_factory=list)
    set_accept_reject: SetAcceptReject


class FunctionalGroupAck(_Base):
    """AK1 + AK9 — one per functional group in the source batch.

    AK901 is implied by the parent AK1; we don't re-emit it.
    AK905 = number of received sets in the source group.
    AK906 = number accepted.
    AK907 = number rejected.
    AK909 = group accept/reject code ("A" accepted, "E" accepted with
    errors, "R" rejected, "P" partially accepted).
    """

    ak1: AcknowledgmentHeader
    received_count: int
    accepted_count: int
    rejected_count: int
    ack_code: Literal["A", "E", "R", "P"]


class ParseResult999(_Base):
    """Top-level parsed 999 ACK document.

    A 999 can carry multiple functional groups (one per GS/GE envelope).
    In v1 Cyclone emits a single 999 per inbound 837, so
    ``functional_group_acks`` is usually length 1; the parser still
    supports multiple groups for forward compatibility.
    """

    envelope: Envelope
    functional_group_acks: list[FunctionalGroupAck] = Field(default_factory=list)
    set_responses: list[SetFunctionalGroupResponse] = Field(default_factory=list)
    summary: BatchSummary


__all__ = [
    "AcknowledgmentHeader",
    "FunctionalGroupAck",
    "ParseResult999",
    "SegmentContext",
    "SegmentError",
    "SetAcceptReject",
    "SetFunctionalGroupResponse",
]


# Silence the "imported but unused" complaint on a few names that exist
# for parity with the 837P / 835 module's public surface but are not
# referenced in this file's implementation.
_ = (Address, BaseModel, ValidationIssue, ValidationReport)
