"""Lab 01 — A2A delegation and a protocol-agnostic core.

MCP answers "what tools do I have". A2A answers "who else can do this, and how do I hand
it to them" — which brings long-running tasks, artifacts, cancellation, push callbacks,
and an identity chain across organizational boundaries.

The rule this lab exists to demonstrate: the KERNEL speaks an internal task model; A2A
and ACP are EDGE ADAPTERS. Otherwise every protocol revision is a rewrite.

    pytest test_lab.py -v
    LAB_MODULE=solution pytest test_lab.py -v    # the reference, must be green
"""

from __future__ import annotations

from dataclasses import dataclass, field, replace
from enum import Enum
from typing import Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

A2A_VERSION = "0.3"

# ======================================================================================
# 1. The parts model
# ======================================================================================


class PartKind(str, Enum):
    TEXT = "text"
    FILE = "file"
    DATA = "data"


@dataclass(frozen=True)
class Part:
    """One piece of a message or artifact: prose, a file, or STRUCTURED data.

    The data kind is what makes agent-to-agent different from chat: a delegated task
    returns a typed result, not a paragraph.
    """

    kind: PartKind
    text: Optional[str] = None
    uri: Optional[str] = None
    mime_type: Optional[str] = None
    data: Optional[Mapping[str, object]] = None

    def __post_init__(self) -> None:
        # TODO: a TEXT part needs text; a FILE part needs a uri; a DATA part needs data.
        #       Anything else -> ValueError.
        raise NotImplementedError

    @classmethod
    def text_part(cls, text: str) -> "Part":
        # TODO
        raise NotImplementedError

    @classmethod
    def file_part(cls, uri: str, mime_type: str) -> "Part":
        # TODO
        raise NotImplementedError

    @classmethod
    def data_part(cls, data: Mapping[str, object]) -> "Part":
        """COPY the mapping — an aliased dict is a mutable frozen dataclass."""
        # TODO
        raise NotImplementedError


class Role(str, Enum):
    USER = "user"      # the CALLER — for a delegated task, another agent
    AGENT = "agent"


@dataclass(frozen=True)
class Message:
    message_id: str
    role: Role
    parts: Tuple[Part, ...]
    task_id: Optional[str] = None
    context_id: Optional[str] = None

    def text(self) -> str:
        """Space-join the text of TEXT parts only."""
        # TODO
        raise NotImplementedError


@dataclass(frozen=True)
class Artifact:
    """A durable OUTPUT, distinct from the conversation about it. The audit record wants
    the artifact, not the chat."""

    artifact_id: str
    name: str
    parts: Tuple[Part, ...]
    description: str = ""


# ======================================================================================
# 2. The task lifecycle
# ======================================================================================


class TaskState(str, Enum):
    SUBMITTED = "submitted"
    WORKING = "working"
    INPUT_REQUIRED = "input-required"
    AUTH_REQUIRED = "auth-required"
    COMPLETED = "completed"
    CANCELED = "canceled"
    FAILED = "failed"
    REJECTED = "rejected"


TERMINAL_STATES: frozenset = frozenset({
    TaskState.COMPLETED, TaskState.CANCELED, TaskState.FAILED, TaskState.REJECTED
})

# TODO: fill the transition map (state -> frozenset of legal next states).
#   SUBMITTED      -> WORKING, REJECTED, CANCELED, FAILED, AUTH_REQUIRED
#   WORKING        -> INPUT_REQUIRED, AUTH_REQUIRED, COMPLETED, CANCELED, FAILED
#   INPUT_REQUIRED -> WORKING, CANCELED, FAILED          (NOT rejected — see the tests)
#   AUTH_REQUIRED  -> WORKING, CANCELED, FAILED, REJECTED
# Note SUBMITTED cannot go straight to COMPLETED: work must be observed to have happened.
TASK_TRANSITIONS: Mapping[TaskState, frozenset] = {}


class IllegalTaskTransition(RuntimeError):
    pass


def advance(state: TaskState, target: TaskState) -> TaskState:
    """Terminal states accept nothing; undeclared edges raise."""
    # TODO
    raise NotImplementedError


@dataclass(frozen=True)
class TaskStatus:
    state: TaskState
    message: Optional[Message] = None
    timestamp: int = 0


@dataclass(frozen=True)
class Task:
    """The unit of delegated work. ``context_id`` groups related tasks into one story."""

    task_id: str
    context_id: str
    status: TaskStatus
    history: Tuple[Message, ...] = ()
    artifacts: Tuple[Artifact, ...] = ()
    delegation_chain: Tuple[str, ...] = ()

    @property
    def is_terminal(self) -> bool:
        return self.status.state in TERMINAL_STATES


# ======================================================================================
# 3. Agent Cards and discovery
# ======================================================================================


@dataclass(frozen=True)
class Skill:
    skill_id: str
    name: str
    description: str
    tags: Tuple[str, ...] = ()
    examples: Tuple[str, ...] = ()


@dataclass(frozen=True)
class AgentCard:
    name: str
    description: str
    url: str
    version: str
    protocol_version: str = A2A_VERSION
    skills: Tuple[Skill, ...] = ()
    streaming: bool = False
    push_notifications: bool = False
    default_input_modes: Tuple[str, ...] = ("text/plain",)
    default_output_modes: Tuple[str, ...] = ("text/plain",)
    security_schemes: Tuple[str, ...] = ("oauth2",)
    owner: str = "unknown"
    tenants: Tuple[str, ...] = ()
    max_data_classification: str = "internal"

    def to_json(self) -> Dict[str, object]:
        """The PUBLISHED card. Exactly these keys:

        protocolVersion, name, description, url, version,
        capabilities{streaming, pushNotifications},
        defaultInputModes, defaultOutputModes, securitySchemes,
        skills[{id, name, description, tags, examples}]

        ``tenants``, ``owner`` and ``max_data_classification`` are PLATFORM metadata and
        must NOT be published.
        """
        # TODO
        raise NotImplementedError


_CLASSIFICATION_ORDER = ("public", "internal", "confidential", "restricted")


def classification_rank(name: str) -> int:
    """Unknown -> ValueError."""
    # TODO
    raise NotImplementedError


@dataclass(frozen=True)
class CallerContext:
    agent_id: str
    tenant: str
    user_id: str
    scopes: Tuple[str, ...] = ()
    data_classification: str = "internal"
    delegation_chain: Tuple[str, ...] = ()


class AgentDirectory:
    """Discovery over Agent Cards. Same rule as Phase 02: filter BEFORE you list."""

    def __init__(self) -> None:
        self._cards: Dict[str, AgentCard] = {}

    def register(self, card: AgentCard) -> None:
        """Duplicate name -> ValueError."""
        # TODO
        raise NotImplementedError

    def get(self, name: str) -> Optional[AgentCard]:
        return self._cards.get(name)

    def discover(self, caller: CallerContext, *, tags: Sequence[str] = ()) -> List[AgentCard]:
        """Cards this caller may see, ranked by best skill-tag overlap desc, then name asc.

        Exclude: the caller itself; cards restricted to other tenants; cards whose
        ``max_data_classification`` is BELOW the caller's ``data_classification`` (they
        cannot handle our data); and, when ``tags`` is non-empty, cards with zero overlap.
        """
        # TODO
        raise NotImplementedError


# ======================================================================================
# 4. Delegation admission — the part A2A does not specify
# ======================================================================================

MAX_DELEGATION_DEPTH = 4


@dataclass(frozen=True)
class DelegationDenial:
    code: str
    reason: str


def check_delegation(
    caller: CallerContext,
    card: AgentCard,
    *,
    max_depth: int = MAX_DELEGATION_DEPTH,
) -> List[DelegationDenial]:
    """Return ALL denials. A2A says how to send a task, not whether you may.

    - DEPTH_EXCEEDED          : len(chain) >= max_depth
    - CYCLE_DETECTED          : card.name already in the chain
    - TENANT_NOT_PERMITTED    : card restricted to other tenants
    - CLASSIFICATION_EXCEEDED : card handles less than the task carries
    - NO_ACCEPTABLE_AUTH      : card offers neither "oauth2" nor "mtls"
    """
    # TODO
    raise NotImplementedError


# ======================================================================================
# 5. Streaming events
# ======================================================================================


@dataclass(frozen=True)
class TaskStatusUpdate:
    task_id: str
    context_id: str
    status: TaskStatus
    final: bool = False


@dataclass(frozen=True)
class TaskArtifactUpdate:
    task_id: str
    context_id: str
    artifact: Artifact
    append: bool = False
    last_chunk: bool = False


StreamEvent = object   # TaskStatusUpdate | TaskArtifactUpdate | Task


@dataclass(frozen=True)
class PushNotificationConfig:
    """A CALLER-supplied URL your server will request. SSRF by construction unless the
    host is allow-listed and the callback carries a verifiable token."""

    url: str
    token: str = ""
    authentication: str = "bearer"


# ======================================================================================
# 6. The A2A server
# ======================================================================================


AgentHandler = Callable[["A2AServer", Task, Message], Iterable[StreamEvent]]


class A2AError(Exception):
    def __init__(self, code: str, message: str) -> None:
        super().__init__(message)
        self.code = code
        self.message = message


class A2AServer:
    def __init__(self, *, card: AgentCard, handler: AgentHandler,
                 allowed_callback_hosts: Sequence[str] = ()) -> None:
        self.card = card
        self.handler = handler
        self.allowed_callback_hosts = tuple(allowed_callback_hosts)
        self.tasks: Dict[str, Task] = {}
        self.push_configs: Dict[str, PushNotificationConfig] = {}
        self.pushed: List[Tuple[str, TaskStatus]] = []
        self._tick = 0
        self._counter = 0

    def _next_tick(self) -> int:
        self._tick += 1
        return self._tick

    def _next_id(self, prefix: str) -> str:
        self._counter += 1
        return f"{prefix}-{self._counter}"

    def _store(self, task: Task) -> Task:
        self.tasks[task.task_id] = task
        return task

    def set_status(self, task: Task, state: TaskState, message: Optional[Message] = None) -> Task:
        """Advance the lifecycle, append the message to history, store, and — if a push
        config exists for this task — append ``(url, status)`` to ``self.pushed``."""
        # TODO
        raise NotImplementedError

    def add_artifact(self, task: Task, artifact: Artifact) -> Task:
        # TODO
        raise NotImplementedError

    def message_send(self, message: Message, *, caller: CallerContext,
                     context_id: Optional[str] = None) -> Task:
        """Blocking form: drain ``message_stream`` and return the final Task."""
        # TODO
        raise NotImplementedError

    def message_stream(self, message: Message, *, caller: CallerContext,
                       context_id: Optional[str] = None) -> Iterable[StreamEvent]:
        """Streaming form. A generator that yields updates and, LAST, the final Task.

        1. ``check_delegation`` first — a denied delegation must NOT create a task.
           Raise A2AError with the first denial's code.
        2. If ``message.task_id`` is set: load it (missing -> TASK_NOT_FOUND, terminal ->
           TASK_TERMINAL) and append the message to history.
           Otherwise: create a Task in SUBMITTED, with
           ``delegation_chain = caller.delegation_chain + (caller.agent_id,)``
           — the chain comes from the CALLER's context, never from the message — and
           yield an initial TaskStatusUpdate.
        3. Delegate to ``self.handler(self, stored_task, message)``, yielding its events.
        4. Yield the stored Task.
        """
        # TODO
        raise NotImplementedError

    def tasks_get(self, task_id: str) -> Task:
        """Missing -> A2AError("TASK_NOT_FOUND", ...)."""
        # TODO
        raise NotImplementedError

    def tasks_cancel(self, task_id: str) -> Task:
        """Already terminal -> A2AError("TASK_TERMINAL", ...)."""
        # TODO
        raise NotImplementedError

    def set_push_config(self, task_id: str, config: PushNotificationConfig) -> PushNotificationConfig:
        """Register a callback.

        - card does not support push -> UNSUPPORTED
        - unknown task -> TASK_NOT_FOUND (via tasks_get)
        - host not in the allow-list (when one is configured) -> CALLBACK_NOT_ALLOWED
        - empty token -> CALLBACK_UNAUTHENTICATED
        """
        # TODO
        raise NotImplementedError


def _host_of(url: str) -> str:
    """Host portion of a URL, without scheme, path or userinfo."""
    # TODO
    raise NotImplementedError


# ======================================================================================
# 7. The A2A client
# ======================================================================================


class A2AClient:
    """The delegating side. Owns the identity chain — the callee cannot be trusted to."""

    def __init__(self, server: A2AServer) -> None:
        self.server = server
        self._counter = 0

    def _next_message_id(self) -> str:
        self._counter += 1
        return f"msg-{self._counter}"

    def card(self) -> AgentCard:
        return self.server.card

    def delegate(self, caller: CallerContext, text: str, *,
                 data: Optional[Mapping[str, object]] = None,
                 context_id: Optional[str] = None) -> Task:
        """Build a USER message with a text part (plus a data part if given) and send."""
        # TODO
        raise NotImplementedError

    def delegate_streaming(self, caller: CallerContext, text: str,
                           context_id: Optional[str] = None) -> List[StreamEvent]:
        # TODO
        raise NotImplementedError

    def reply(self, caller: CallerContext, task: Task, text: str) -> Task:
        """Continue an existing task — set ``task_id`` and ``context_id`` on the message."""
        # TODO
        raise NotImplementedError


# ======================================================================================
# 8. The protocol-agnostic core, and an ACP adapter
# ======================================================================================


@dataclass(frozen=True)
class InternalTask:
    """What the kernel stores — deliberately neither A2A's shape nor ACP's."""

    task_id: str
    context_id: str
    state: str
    input_text: str
    output_text: str = ""
    structured: Mapping[str, object] = field(default_factory=dict)
    delegation_chain: Tuple[str, ...] = ()
    artifacts: Tuple[str, ...] = ()


# TODO: map every A2A state onto our vocabulary. Suggested:
#   submitted -> queued        working -> running       input-required -> awaiting_input
#   auth-required -> awaiting_auth                      completed -> succeeded
#   canceled -> cancelled      failed -> failed         rejected -> rejected
A2A_STATE_TO_INTERNAL: Mapping[TaskState, str] = {}

# TODO: the inverse. It must be a true inverse — a test round-trips every state.
INTERNAL_TO_A2A_STATE: Mapping[str, TaskState] = {}

# TODO: ACP's REST vocabulary onto ours.
#   created -> queued   in-progress -> running   awaiting -> awaiting_input
#   completed -> succeeded   cancelled -> cancelled   failed -> failed
ACP_STATUS_TO_INTERNAL: Mapping[str, str] = {}

# TODO: ours onto ACP's. NOT injective: ACP has no "rejected" and no separate auth state,
# so rejected -> "failed" and awaiting_auth -> "awaiting". Lossy mappings are fine as long
# as they are DECLARED — a test pins both.
INTERNAL_TO_ACP_STATUS: Mapping[str, str] = {}


def a2a_to_internal(task: Task) -> InternalTask:
    """Flatten a Task: text parts of all artifacts joined into ``output_text``, data
    parts merged into ``structured``, the FIRST user message as ``input_text``, artifact
    NAMES into ``artifacts``."""
    # TODO
    raise NotImplementedError


def internal_to_acp(task: InternalTask) -> Dict[str, object]:
    """An ACP-shaped run envelope:

        {run_id, session_id, status, input:[{content_type, content}],
         output:[{role:"agent", parts:[...]}],
         metadata:{delegation_chain, artifacts}}

    Emit a text part only when ``output_text`` is non-empty, and a JSON part only when
    ``structured`` is non-empty — so an empty task round-trips.
    """
    # TODO
    raise NotImplementedError


def acp_to_internal(envelope: Mapping[str, object]) -> InternalTask:
    """The inverse. Unknown status -> ValueError."""
    # TODO
    raise NotImplementedError


def main() -> None:
    print("implement the TODOs, then compare with `python solution.py`")


if __name__ == "__main__":
    main()