"""Language-neutral binary payload handles used by canonical stream contracts."""

from __future__ import annotations

import base64
import binascii
import hashlib
from typing import Annotated, Literal

from pydantic import Field, NonNegativeInt, PositiveInt, model_validator

from .common import ContractModel, Identifier, Sha256Digest, ShortText


class InlinePayload(ContractModel):
    kind: Literal["inline-base64"] = "inline-base64"
    media_type: ShortText
    byte_length: NonNegativeInt
    sha256: Sha256Digest
    data: str

    @model_validator(mode="after")
    def inline_data_must_match_descriptor(self) -> InlinePayload:
        try:
            decoded = base64.b64decode(self.data, validate=True)
        except (binascii.Error, ValueError) as exc:
            raise ValueError("data must be canonical base64") from exc
        if len(decoded) != self.byte_length:
            raise ValueError("byte_length does not match decoded inline payload")
        digest = f"sha256:{hashlib.sha256(decoded).hexdigest()}"
        if digest != self.sha256:
            raise ValueError("sha256 does not match decoded inline payload")
        if base64.b64encode(decoded).decode("ascii") != self.data:
            raise ValueError("data must use canonical padded base64 encoding")
        return self


class ReferencedPayload(ContractModel):
    """Opaque payload location owned by a declared transport or evidence store."""

    kind: Literal["reference"] = "reference"
    media_type: ShortText
    byte_length: NonNegativeInt
    sha256: Sha256Digest | None = None
    reference_scheme: Literal["evidence", "shared-memory", "transport", "uri"]
    locator: ShortText
    offset: NonNegativeInt = 0
    segment_length: PositiveInt | None = None
    owner_id: Identifier | None = None

    @model_validator(mode="after")
    def segment_must_fit_known_payload(self) -> ReferencedPayload:
        if self.offset > self.byte_length:
            raise ValueError("offset cannot exceed byte_length")
        if self.segment_length is not None and self.offset + self.segment_length > self.byte_length:
            raise ValueError("referenced segment exceeds byte_length")
        if self.reference_scheme in {"evidence", "shared-memory"} and self.owner_id is None:
            raise ValueError("evidence and shared-memory references require owner_id")
        return self


PayloadHandle = Annotated[InlinePayload | ReferencedPayload, Field(discriminator="kind")]
