"""Centralized A2A protocol-version handling.

Every place the CLI has to know *which* revision of the A2A wire format it is
speaking lives here, so moving to a future revision (1.1 is expected upstream)
means editing this module rather than every A2A call site.

The CLI speaks **A2A 1.0**. A2A 0.3 is retained only as an emergency rollback
lever, selected through the ``WORKIQ_A2A_PROTOCOL_VERSION`` environment
variable; nothing ever downgrades automatically. The differences the profiles
below encapsulate:

===================== ============================ ==========================
Concern               0.3                          1.0
===================== ============================ ==========================
Send method           ``message/send``             ``SendMessage``
Message role          ``user`` / ``agent``         ``ROLE_USER`` / ``ROLE_AGENT``
Task state            ``completed``                ``TASK_STATE_COMPLETED``
Part discrimination   ``kind`` field               member presence
Part media type       nested / absent              top-level ``mediaType``
Send result           bare object with ``kind``    ``{"task": …}`` / ``{"message": …}``
Blocking send         ``blocking: true``           ``returnImmediately: false``
===================== ============================ ==========================

**Only writers are versioned.** The readers below are deliberately tolerant of
both wire shapes — member-based *and* ``kind``-based parts, wrapped *and* bare
send results, prefixed *and* bare enum values — regardless of the active
profile, so a response is always parsed on its own merits. That is what keeps
the rollback lever cheap: it changes four fields on the way out and nothing on
the way in.

**Maintenance rule: 0.3 is frozen.** It is preserved exactly as it behaved
before the 1.0 upgrade, and no new capability should grow a 0.3 code path. It
is deliberately not exposed as a command-line flag, because a documented flag
is far harder to retract than an environment variable once the 1.0 rollout is
proven. Without that rule every future A2A feature pays a dual-stack cost.
"""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Optional, Tuple

from error_messages import (
    ERROR_CODE_AGENT_REQUEST_FAILED,
    AgentRequestError,
)

PROTOCOL_VERSION_1_0 = "1.0"
PROTOCOL_VERSION_0_3 = "0.3"

# What the CLI speaks unless the rollback lever is pulled.
DEFAULT_PROTOCOL_VERSION = PROTOCOL_VERSION_1_0

# Every revision the client can be pointed at. Ordered for stable error
# messages, newest first.
SELECTABLE_PROTOCOL_VERSIONS: Tuple[str, ...] = (
    PROTOCOL_VERSION_1_0,
    PROTOCOL_VERSION_0_3,
)

#: Rollback lever. Undocumented in ``--help`` on purpose: this exists for an
#: incident, not for routine configuration.
ENV_PROTOCOL_VERSION = "WORKIQ_A2A_PROTOCOL_VERSION"

# Sent on every request so a server doing version negotiation answers in kind.
A2A_VERSION_HEADER = "A2A-Version"

# JSON-RPC "method not found". Reachable from this client even though our
# profiles always pair a version with that version's method name: the server
# picks its method table from the ``A2A-Version`` header and treats an
# absent-or-empty header as 0.3, so anything that strips or blanks that header
# in transit (a proxy, a gateway) turns our valid 1.0 ``SendMessage`` into this
# error. Verified against the live endpoint. A genuinely unknown method also
# lands here, so the wording it produces hedges rather than asserting.
JSONRPC_METHOD_NOT_FOUND = -32601

# A2A ``VersionNotSupportedError`` (spec 3.6.2): the ``A2A-Version`` header
# named a revision the endpoint does not serve. Unlike the above this is
# unambiguous -- the header arrived and was rejected on its merits -- and the
# server's message names the versions it does support, so it is always worth
# surfacing verbatim. (An unrecognized value yields this; an empty one does
# not, because empty is read as "no header" and falls back to 0.3.)
JSONRPC_VERSION_NOT_SUPPORTED = -32009

#: Error codes that mean "this endpoint will not serve the version we asked
#: for", so a send rejected with one of these deserves a message naming the
#: mismatch rather than a generic protocol error.
PROTOCOL_VERSION_ERROR_CODES: Tuple[int, ...] = (
    JSONRPC_VERSION_NOT_SUPPORTED,
    JSONRPC_METHOD_NOT_FOUND,
)

# Prefixes 1.0 adds to enum values, stripped when normalizing for comparison.
_TASK_STATE_PREFIX = "task_state_"
_ROLE_PREFIX = "role_"

# Result payload keys of the 1.0 ``SendMessageResponse`` oneof.
_RESULT_KIND_TASK = "task"
_RESULT_KIND_MESSAGE = "message"


class A2AProtocolVersionError(AgentRequestError):
    """Raised when an endpoint cannot serve the protocol version we require.

    Carries the required version so the message can name the mismatch rather
    than failing with a generic protocol error further down the call.
    """

    def __init__(
        self,
        message: str,
        *,
        required: str = DEFAULT_PROTOCOL_VERSION,
    ) -> None:
        super().__init__(message, code=ERROR_CODE_AGENT_REQUEST_FAILED)
        self.required = required


# ---------------------------------------------------------------------- #
#  Version profiles (the only version-dependent code)                     #
# ---------------------------------------------------------------------- #


@dataclass(frozen=True)
class ProtocolProfile:
    """The wire dialect for one A2A revision.

    Holds every outbound difference between revisions, so adding a future
    revision means adding one profile rather than branching at call sites.
    """

    version: str
    send_method: str
    role_user: str
    #: 0.3 tags messages and parts with a ``kind`` discriminator; 1.0 uses
    #: member presence instead.
    uses_kind_discriminator: bool

    @property
    def is_default(self) -> bool:
        """Whether this is the revision the CLI speaks by default."""
        return self.version == DEFAULT_PROTOCOL_VERSION

    def text_part(self, text: str) -> Dict[str, Any]:
        """Build a text ``Part`` in this revision's shape."""
        if self.uses_kind_discriminator:
            return {"kind": "text", "text": text}
        return {"text": text}

    def data_part(
        self,
        data: Dict[str, Any],
        media_type: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Build a data ``Part`` in this revision's shape.

        On 1.0 ``mediaType`` is emitted as a **top-level sibling** of ``data``.
        On 0.3 it is dropped, because that revision's ``Part`` cannot express
        it.
        """
        if self.uses_kind_discriminator:
            return {"kind": "data", "data": data}
        part: Dict[str, Any] = {"data": data}
        if media_type:
            part["mediaType"] = media_type
        return part

    def message_envelope(self, parts: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Build the version-appropriate skeleton of an outbound message."""
        message: Dict[str, Any] = {"role": self.role_user, "parts": parts}
        if self.uses_kind_discriminator:
            message["kind"] = "message"
        return message

    def send_configuration(self, modes: List[str]) -> Dict[str, Any]:
        """Build the synchronous-send ``configuration`` block.

        1.0 replaced 0.3's ``blocking: true`` with the inverted
        ``returnImmediately``.
        """
        configuration: Dict[str, Any] = {"acceptedOutputModes": modes}
        if self.uses_kind_discriminator:
            configuration["blocking"] = True
        else:
            configuration["returnImmediately"] = False
        return configuration


_PROFILE_1_0 = ProtocolProfile(
    version=PROTOCOL_VERSION_1_0,
    send_method="SendMessage",
    role_user="ROLE_USER",
    uses_kind_discriminator=False,
)

_PROFILE_0_3 = ProtocolProfile(
    version=PROTOCOL_VERSION_0_3,
    send_method="message/send",
    role_user="user",
    uses_kind_discriminator=True,
)

_PROFILES: Dict[str, ProtocolProfile] = {
    PROTOCOL_VERSION_1_0: _PROFILE_1_0,
    PROTOCOL_VERSION_0_3: _PROFILE_0_3,
}


def get_profile(version: Optional[str] = None) -> ProtocolProfile:
    """Return the profile for ``version``, defaulting to the current revision.

    Raises:
        A2AProtocolVersionError: When ``version`` names a revision this client
            has no profile for.
    """
    if version is None:
        return _PROFILES[DEFAULT_PROTOCOL_VERSION]
    normalized = normalize_protocol_version(version)
    profile = _PROFILES.get(normalized)
    if profile is None:
        raise A2AProtocolVersionError(
            f"Unsupported A2A protocol version {version!r}. "
            f"Supported versions: {', '.join(SELECTABLE_PROTOCOL_VERSIONS)}."
        )
    return profile


def resolve_protocol_version() -> str:
    """Resolve the protocol version to speak from the environment.

    Returns :data:`DEFAULT_PROTOCOL_VERSION` unless
    ``WORKIQ_A2A_PROTOCOL_VERSION`` names another supported revision. An
    unrecognized value falls back to the default rather than raising, because
    an incident-time typo should not take the run down — but it is warned
    about, since silently ignoring the lever would leave an operator believing
    a rollback is in effect when it is not.
    """
    raw = os.environ.get(ENV_PROTOCOL_VERSION)
    if raw is None or not raw.strip():
        return DEFAULT_PROTOCOL_VERSION
    normalized = normalize_protocol_version(raw)
    if normalized in _PROFILES:
        return normalized
    _warn_invalid_protocol_version(raw)
    return DEFAULT_PROTOCOL_VERSION


def _warn_invalid_protocol_version(raw: str) -> None:
    # Imported lazily to avoid a circular import at module load time
    # (cli_logging is a higher-level package than this module, which is
    # imported by common).
    from cli_logging.cli_logger import emit_structured_log
    from cli_logging.logging_utils import Operation

    emit_structured_log(
        "warning",
        f"Ignoring invalid {ENV_PROTOCOL_VERSION}={raw!r}; "
        f"expected one of {', '.join(SELECTABLE_PROTOCOL_VERSIONS)}. "
        f"Using default {DEFAULT_PROTOCOL_VERSION}.",
        Operation.SETUP,
    )


# ---------------------------------------------------------------------- #
#  Version strings                                                        #
# ---------------------------------------------------------------------- #


def normalize_protocol_version(raw: Any) -> str:
    """Reduce a protocol-version string to its ``major.minor`` form.

    ``"1.0.0"``, ``"1.0"`` and ``" v1.0 "`` all normalize to ``"1.0"``.
    Returns ``""`` when the value is missing or not a string.
    """
    if not isinstance(raw, str):
        return ""
    value = raw.strip().lstrip("vV")
    if not value:
        return ""
    parts = value.split(".")
    if len(parts) >= 2:
        return f"{parts[0]}.{parts[1]}"
    return parts[0]


# ---------------------------------------------------------------------- #
#  Enum values (version-tolerant readers)                                 #
# ---------------------------------------------------------------------- #


def normalize_enum_value(value: Any, prefix: str = "") -> str:
    """Fold a 1.0 or 0.3 enum value into lowercase underscore form.

    Handles the hyphenated 0.3 spelling, the ``SCREAMING_SNAKE_CASE`` 1.0
    spelling, and the type prefix 1.0 adds. ``"input-required"``,
    ``"TASK_STATE_INPUT_REQUIRED"`` and ``"input_required"`` all normalize to
    ``"input_required"`` when ``prefix`` is ``"task_state_"``.
    """
    if not isinstance(value, str):
        return ""
    normalized = value.strip().lower().replace("-", "_")
    if prefix and normalized.startswith(prefix):
        normalized = normalized[len(prefix) :]
    return normalized


def normalize_task_state(state: Any) -> str:
    """Fold an A2A task state into lowercase underscore form without prefix."""
    return normalize_enum_value(state, _TASK_STATE_PREFIX)


def normalize_role(role: Any) -> str:
    """Fold an A2A message role into lowercase underscore form without prefix."""
    return normalize_enum_value(role, _ROLE_PREFIX)


# ---------------------------------------------------------------------- #
#  Parts (version-tolerant readers)                                       #
# ---------------------------------------------------------------------- #


def part_text(part: Any) -> Optional[str]:
    """Return the text of a ``Part``, or ``None`` if it is not a text part.

    Accepts both the 1.0 member-based shape (``{"text": …}``) and the 0.3
    shape (``{"kind": "text", "text": …}``).
    """
    if not isinstance(part, dict):
        return None
    value = part.get("text")
    return value if isinstance(value, str) else None


def part_data(part: Any) -> Optional[Dict[str, Any]]:
    """Return the data payload of a ``Part``, or ``None`` if it carries none."""
    if not isinstance(part, dict):
        return None
    value = part.get("data")
    return value if isinstance(value, dict) else None


def join_text_parts(parts: Optional[Iterable[Any]]) -> str:
    """Concatenate every text part in ``parts``, newline-separated.

    Mirrors the 0.3 client's behaviour exactly — including keeping empty
    strings — so migrating the discriminator does not change extracted text.
    """
    values: List[str] = []
    for part in parts or ():
        text = part_text(part)
        if text is not None:
            values.append(text)
    return "\n".join(values)


# ---------------------------------------------------------------------- #
#  Send result envelope (version-tolerant reader)                         #
# ---------------------------------------------------------------------- #


def unwrap_send_result(result: Any) -> Tuple[str, Dict[str, Any]]:
    """Split a ``SendMessage`` result into ``(payload_kind, payload)``.

    1.0 discriminates the ``SendMessageResponse`` oneof by member name
    (``{"task": …}`` / ``{"message": …}``); 0.3 returned the object inline with
    a ``kind`` field. Both are accepted regardless of the active profile.

    Deliberately does **not** guess from structure when neither discriminator
    is present: a well-formed 1.0 result always carries the member and a
    well-formed 0.3 result always carries ``kind``, so anything else is
    malformed. Guessing there risks reading a payload as the wrong type and
    silently extracting no text — a blank answer scored as a real one — which
    is far worse than failing loudly.

    Returns ``("", {})`` when the shape is recognizable as neither, or when it
    is ambiguously both, leaving the caller to raise a protocol error with its
    own wording.
    """
    if not isinstance(result, dict):
        return "", {}

    members = [
        key
        for key in (_RESULT_KIND_TASK, _RESULT_KIND_MESSAGE)
        if isinstance(result.get(key), dict)
    ]
    if len(members) == 1:
        return members[0], result[members[0]]
    if members:
        return "", {}

    kind = normalize_enum_value(result.get("kind"))
    if kind in (_RESULT_KIND_TASK, _RESULT_KIND_MESSAGE):
        return kind, result
    return "", {}


# ---------------------------------------------------------------------- #
#  Version mismatch reporting                                             #
# ---------------------------------------------------------------------- #


def _downgrade_hint(active: ProtocolProfile) -> str:
    """Point at the rollback lever, but only when it is a step down.

    Names the environment variable rather than a flag because the lever is
    deliberately not part of the documented command-line surface.
    """
    if not active.is_default:
        return ""
    return (
        f" To retry using A2A {PROTOCOL_VERSION_0_3}, set "
        f"{ENV_PROTOCOL_VERSION}={PROTOCOL_VERSION_0_3} in the environment."
    )


def _server_detail(server_message: Optional[str]) -> str:
    """Quote the server's own explanation, which we must not swallow.

    A ``VersionNotSupportedError`` names the versions the endpoint *does*
    serve, which is the single most useful fact for whoever has to react.
    Terminating punctuation is added when absent so the quoted text does not
    run into the sentence that follows it.
    """
    text = (server_message or "").strip()
    if not text:
        return ""
    if text[-1] not in ".!?":
        text += "."
    return f" Server said: {text}"


def protocol_version_error(
    agent_url: str,
    profile: ProtocolProfile,
    *,
    code: int,
    server_message: Optional[str] = None,
) -> A2AProtocolVersionError:
    """Build the error for a send the endpoint refused on version grounds.

    This is the *only* protocol-version check the client performs. The agent
    card is never consulted for version compatibility: a card can advertise a
    stale version while the endpoint already serves a newer one, so refusing to
    send based on metadata risks blocking an endpoint that works. Asking the
    server directly costs one round trip, and both codes handled here are
    pre-dispatch rejections, so nothing was executed.

    Two codes reach this path, and they mean different things:

    ``-32009`` (``VersionNotSupportedError``)
        The endpoint rejected the version outright. Unambiguous.
    ``-32601`` (method not found)
        The endpoint does not know our send method. Because the server picks
        its method table from the ``A2A-Version`` header, this is *usually* a
        version mismatch — but it is also what a genuinely unknown method
        returns, so the wording hedges rather than asserting.
    """
    if code == JSONRPC_VERSION_NOT_SUPPORTED:
        summary = (
            f"The agent endpoint does not support A2A protocol version "
            f"{profile.version}."
        )
    else:
        summary = (
            f"The agent endpoint rejected the A2A {profile.version} method "
            f"'{profile.send_method}' as unknown, which usually means it does "
            f"not serve protocol version {profile.version}."
        )
    return A2AProtocolVersionError(
        f"{summary} Endpoint: {agent_url}."
        f"{_server_detail(server_message)}"
        f"{_downgrade_hint(profile)}",
        required=profile.version,
    )