"""Runtime context handed to skill handlers.

The same agent code runs unchanged on local dev, Docker, Kubernetes, and
hosted runtimes — the runtime provides a concrete :class:`RunContext` that
implements artifact storage, secret access, streaming, and cancellation.
"""
from __future__ import annotations

import asyncio
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Generic, Sequence, TypeVar

from pydantic import BaseModel

from .workspace import WorkspaceClient

AuthT = TypeVar("AuthT", bound=BaseModel)


class CancelledByCaller(RuntimeError):
    """Raised by :meth:`RunContext.check_cancelled` when the caller cancelled."""


class MissingScopes(PermissionError):
    """Raised by :meth:`RunContext.require_scopes` when caller lacks scopes."""

    def __init__(self, missing: Sequence[str]) -> None:
        self.missing = tuple(missing)
        super().__init__(f"missing scopes: {sorted(self.missing)}")


@dataclass(frozen=True)
class ArtifactRef:
    """Opaque handle to a stored artifact (blob, file, etc.)."""

    name: str
    uri: str
    mime_type: str
    size_bytes: int


@dataclass(frozen=True)
class AgentEvent:
    """A structured event emitted during a skill run."""

    kind: str
    payload: dict[str, Any] = field(default_factory=dict)


class RunContext(ABC, Generic[AuthT]):
    """Per-invocation context.

    A new context is constructed by the runtime for every skill call. It
    carries caller identity (``auth``), the task identity, and runtime
    capabilities (artifacts, secrets, streaming, cancellation).

    Agents must depend only on this abstract interface, never on a concrete
    runtime implementation.
    """

    task_id: str
    auth: AuthT

    @abstractmethod
    async def emit_event(self, event: AgentEvent) -> None:
        """Publish a structured event to subscribers (UI, logs, traces)."""

    @abstractmethod
    async def write_artifact(
        self, name: str, data: bytes, mime_type: str
    ) -> ArtifactRef:
        """Persist ``data`` as a named artifact and return a reference."""

    @abstractmethod
    async def check_cancelled(self) -> None:
        """Raise :class:`CancelledByCaller` if the caller cancelled."""

    @abstractmethod
    def secret(self, name: str) -> str:
        """Look up a runtime-injected secret by logical name."""

    @property
    @abstractmethod
    def workspace(self) -> WorkspaceClient:
        """Negotiation surface for workspace access.

        Raises if the agent's :attr:`A2AAgent.workspace_access` is disabled.
        """

    # --- concrete helpers built on emit_event ---

    async def emit_progress(self, message: str) -> None:
        """Emit a human-readable progress event."""
        await self.emit_event(AgentEvent(kind="progress", payload={"message": message}))

    async def emit_text_delta(self, text: str) -> None:
        """Emit a streamed token chunk (for LLM-style streaming output)."""
        await self.emit_event(AgentEvent(kind="text_delta", payload={"text": text}))

    async def emit_artifact(self, ref: ArtifactRef) -> None:
        """Notify subscribers that a new artifact is available."""
        await self.emit_event(
            AgentEvent(
                kind="artifact",
                payload={
                    "name": ref.name,
                    "uri": ref.uri,
                    "mime_type": ref.mime_type,
                    "size_bytes": ref.size_bytes,
                },
            )
        )

    async def emit_error(self, message: str, *, code: str | None = None) -> None:
        """Emit a structured error event (does not raise)."""
        await self.emit_event(
            AgentEvent(kind="error", payload={"message": message, "code": code})
        )

    def require_scopes(self, required: Sequence[str]) -> None:
        """Raise :class:`MissingScopes` if ``self.auth`` lacks any required scope.

        Auth models without a ``scopes`` attribute (e.g. :class:`NoAuth`) are
        treated as having an empty scope set.
        """
        if not required:
            return
        auth_scopes = set(getattr(self.auth, "scopes", ()) or ())
        missing = [s for s in required if s not in auth_scopes]
        if missing:
            raise MissingScopes(missing)


class LocalRunContext(RunContext[AuthT]):
    """In-memory context for local dev and tests.

    Stores events and artifacts in lists/dicts. Secrets come from a plain
    mapping. Cancellation is driven by an :class:`asyncio.Event`.
    """

    def __init__(
        self,
        *,
        auth: AuthT,
        task_id: str = "local-task",
        secrets: dict[str, str] | None = None,
        workspace: WorkspaceClient | None = None,
    ) -> None:
        self.task_id = task_id
        self.auth = auth
        self._secrets: dict[str, str] = dict(secrets or {})
        self._workspace = workspace
        self._cancel = asyncio.Event()
        self.events: list[AgentEvent] = []
        self.artifacts: dict[str, bytes] = {}

    @property
    def workspace(self) -> WorkspaceClient:
        if self._workspace is None:
            raise PermissionError(
                "no workspace bound to this context; agent did not declare "
                "workspace_access or runtime did not provision one"
            )
        return self._workspace

    async def emit_event(self, event: AgentEvent) -> None:
        self.events.append(event)

    async def write_artifact(
        self, name: str, data: bytes, mime_type: str
    ) -> ArtifactRef:
        self.artifacts[name] = data
        return ArtifactRef(
            name=name,
            uri=f"memory://{self.task_id}/{name}",
            mime_type=mime_type,
            size_bytes=len(data),
        )

    async def check_cancelled(self) -> None:
        if self._cancel.is_set():
            raise CancelledByCaller(self.task_id)

    def cancel(self) -> None:
        self._cancel.set()

    def secret(self, name: str) -> str:
        try:
            return self._secrets[name]
        except KeyError as exc:
            raise KeyError(f"unknown secret: {name!r}") from exc