"""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, TypeAdapter

from dataclasses import dataclass as _dataclass

from .a2a_client import A2AClient, CallResult
from .discovery import DiscoveryClient
from .grants import Grant, GrantInvalid, verify_grant
from .sandbox import SandboxClient, SandboxUnavailable
from .workspace import WorkspaceClient


@_dataclass(frozen=True)
class LLMCreds:
    """An LLM endpoint + credentials handed to the skill at runtime.

    Resolution order — checked in :meth:`RunContext.llm`:

      1. ``llm_creds`` in the inbound /invoke body — usually set by the
         control plane when the caller has registered their own keys
         and the agent's Card declares ``llm_provisioning=caller_provided``.
      2. Per-agent env vars: ``AGENT_LLM_{URL,KEY,MODEL}`` for
         ``llm_provisioning=agent_byok``.
      3. Platform defaults: ``A2A_LITELLM_{URL,KEY,MODEL}``.

    Skill code shouldn't care which branch fired — just read
    ``ctx.llm`` and pass it through to your chat client.
    """
    base_url: str
    api_key: str
    model: str
    source: str  # "caller" | "agent_byok" | "platform"

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)}")


class ScopeDenied(PermissionError):
    """Raised by :meth:`RunContext.request_scope` when the platform refuses."""


class ScopeExpansionNotAllowed(PermissionError):
    """Raised when a skill calls request_scope without opting in via @skill."""


@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 _AskRegistry:
    """Process-local registry of pending :meth:`RunContext.ask` calls.

    Each call gets a unique question_id; the runtime layer is responsible
    for routing the answer back via :meth:`answer`. Lives at module scope
    so the HTTP /answers handler can resolve waiters across requests.
    """

    _waiters: dict[str, "asyncio.Future[str]"] = {}

    @classmethod
    async def ask(
        cls, ctx: "RunContext[Any]", prompt: str, *, timeout: float = 180.0
    ) -> str:
        import secrets

        question_id = f"q_{secrets.token_hex(6)}"
        loop = asyncio.get_running_loop()
        fut: asyncio.Future[str] = loop.create_future()
        cls._waiters[question_id] = fut
        try:
            await ctx.emit_event(
                AgentEvent(
                    kind="question",
                    payload={"question_id": question_id, "prompt": prompt},
                )
            )
            return await asyncio.wait_for(fut, timeout=timeout)
        finally:
            cls._waiters.pop(question_id, None)

    @classmethod
    def answer(cls, question_id: str, answer: str) -> bool:
        fut = cls._waiters.get(question_id)
        if fut is None or fut.done():
            return False
        fut.set_result(answer)
        return True


class _InputRegistry:
    """Process-local registry of pending :meth:`RunContext.collect` calls."""

    _waiters: dict[str, "asyncio.Future[dict[str, Any]]"] = {}

    @classmethod
    async def collect(
        cls,
        ctx: "RunContext[Any]",
        *,
        title: str,
        reason: str,
        schema: dict[str, Any],
        ui_schema: dict[str, Any] | None,
        timeout: float,
    ) -> dict[str, Any]:
        import secrets

        request_id = f"ir_{secrets.token_hex(6)}"
        loop = asyncio.get_running_loop()
        fut: asyncio.Future[dict[str, Any]] = loop.create_future()
        cls._waiters[request_id] = fut
        try:
            await ctx.emit_event(
                AgentEvent(
                    kind="input_request",
                    payload={
                        "request_id": request_id,
                        "title": title,
                        "reason": reason,
                        "schema": schema,
                        "ui_schema": ui_schema or {},
                    },
                )
            )
            return await asyncio.wait_for(fut, timeout=timeout)
        finally:
            cls._waiters.pop(request_id, None)

    @classmethod
    def submit(cls, request_id: str, value: dict[str, Any]) -> bool:
        fut = cls._waiters.get(request_id)
        if fut is None or fut.done():
            return False
        fut.set_result(value)
        return True


class _ScopeRegistry:
    """Process-local registry of pending :meth:`RunContext.request_scope` calls.

    Each call gets a unique request_id; the runtime layer routes the
    platform's reply back via :meth:`resolve` (delivers a fresh signed
    grant token) or :meth:`deny` (delivers a deny reason).
    """

    _waiters: dict[str, "asyncio.Future[Grant | str]"] = {}

    @classmethod
    async def request(
        cls,
        ctx: "RunContext[Any]",
        *,
        reason: str,
        read: Sequence[str],
        write_prefix: str | None,
        ttl_seconds: int,
        mode: str,
        timeout: float,
    ) -> Grant:
        import secrets

        request_id = f"sr_{secrets.token_hex(6)}"
        loop = asyncio.get_running_loop()
        fut: asyncio.Future[Grant | str] = loop.create_future()
        cls._waiters[request_id] = fut
        try:
            await ctx.emit_event(
                AgentEvent(
                    kind="scope_request",
                    payload={
                        "request_id": request_id,
                        "reason": reason,
                        "read_patterns": list(read),
                        "write_prefix": write_prefix,
                        "ttl_seconds": ttl_seconds,
                        "mode": mode,
                    },
                )
            )
            outcome = await asyncio.wait_for(fut, timeout=timeout)
        finally:
            cls._waiters.pop(request_id, None)
        if isinstance(outcome, str):
            raise ScopeDenied(outcome)
        return outcome

    @classmethod
    def resolve(cls, request_id: str, grant_token: str) -> bool:
        fut = cls._waiters.get(request_id)
        if fut is None or fut.done():
            return False
        try:
            grant = verify_grant(grant_token)
        except GrantInvalid as exc:
            fut.set_result(f"invalid grant returned by platform: {exc}")
            return True
        fut.set_result(grant)
        return True

    @classmethod
    def deny(cls, request_id: str, reason: str) -> bool:
        fut = cls._waiters.get(request_id)
        if fut is None or fut.done():
            return False
        fut.set_result(reason)
        return True


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
    def llm(self) -> LLMCreds:
        """LLM endpoint + credentials for this invocation.

        Skill code that wants to make an LLM call should read this
        instead of hard-coding env-var lookups. The platform decides
        which creds to hand over based on the agent's
        :class:`LLMProvisioning` declaration + per-request data the
        control plane forwards (caller's keys, etc.).
        """
        import os

        forwarded = getattr(self, "_llm_creds", None)
        if isinstance(forwarded, LLMCreds):
            return forwarded
        # Agent-BYOK overrides — per-agent env vars take precedence over
        # the platform defaults so authors can swap providers without
        # rebuilding the CP.
        if os.environ.get("AGENT_LLM_KEY"):
            return LLMCreds(
                base_url=os.environ.get("AGENT_LLM_URL", "https://api.openai.com/v1"),
                api_key=os.environ["AGENT_LLM_KEY"],
                model=os.environ.get("AGENT_LLM_MODEL", "gpt-4o"),
                source="agent_byok",
            )
        return LLMCreds(
            base_url=(
                os.environ.get("A2A_LITELLM_URL",
                               "http://litellm.llm.svc.cluster.local:4000") + "/v1"
            ),
            api_key=os.environ.get("A2A_LITELLM_KEY", "sk-microcash-local"),
            model=os.environ.get("A2A_LITELLM_MODEL", "gpt-5.5"),
            source="platform",
        )

    def workspace_backend(self, *, image: str = "python:3.11-slim") -> Any:
        """Return a durable backend bound to this invocation workspace.

        Pass this to ``deepagents.create_deep_agent(..., backend=...)`` so
        DeepAgents' built-in file tools write to the caller's durable
        workspace instead of the default ephemeral LangGraph state backend.
        When the runtime attached ``ctx.sandbox``, the backend's ``execute``
        tool also runs in a sandbox with the same workspace mounted at
        ``/workspace``.
        """
        from .deepagents import WorkspaceBackend

        try:
            sandbox = self.sandbox
        except Exception:  # noqa: BLE001
            sandbox = None
        return WorkspaceBackend(
            self.workspace,
            sandbox=sandbox,
            default_image=image,
        )

    def deepagents_backend(self, *, image: str = "python:3.11-slim") -> Any:
        """Compatibility alias for :meth:`workspace_backend`."""
        return self.workspace_backend(image=image)

    @property
    def cp_jwt(self) -> str | None:
        """Caller's control-plane JWT, when the platform forwarded it.

        Populated only if the agent class declares ``wants_cp_jwt=True``;
        otherwise always ``None``. Skills that need to call back into
        ``/v1/me/*`` on the user's behalf read this and pass it as the
        ``authorization: bearer ...`` header.
        """
        return getattr(self, "_cp_jwt", None)

    @property
    def cp_url(self) -> str | None:
        """Control-plane base URL paired with :attr:`cp_jwt`."""
        return getattr(self, "_cp_url", None)

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

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

    @property
    @abstractmethod
    def sandbox(self) -> SandboxClient:
        """Code-execution surface (microsandbox-backed by default).

        Raises :class:`SandboxUnavailable` if the runtime did not attach a
        sandbox client to this context (e.g. local dev with no host daemon).
        """

    @property
    @abstractmethod
    def discover(self) -> DiscoveryClient:
        """Registry-backed discovery: find other agents by tag/capability/skill."""

    async def call(
        self,
        target: str,
        skill: str,
        *,
        args: dict[str, Any] | None = None,
        grant: str | None = None,
        timeout: float | None = None,
    ) -> CallResult:
        """Invoke another agent's skill via the runtime's :class:`A2AClient`.

        ``target`` is whatever the underlying client expects — an HTTP URL
        for :class:`HttpA2AClient`, an agent name for in-process routing.
        Pair with :meth:`WorkspaceClient.delegate` to hand a scoped
        workspace grant to the callee.
        """
        client = self._a2a_client()
        return await client.call(target, skill, args=args, grant=grant, timeout=timeout)

    @abstractmethod
    def _a2a_client(self) -> A2AClient:
        """Return the runtime's outbound A2A client (or raise if absent)."""

    # --- 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 ask(self, prompt: str, *, timeout: float = 180.0) -> str:
        """Pause the skill until the caller answers a free-text question.

        Emits an event with ``kind="question"``. The runtime is responsible
        for routing the answer back via :meth:`answer`. If no answer arrives
        within ``timeout`` seconds, raises :class:`asyncio.TimeoutError`.
        """
        return await _AskRegistry.ask(self, prompt, timeout=timeout)

    async def collect(
        self,
        schema: type[BaseModel] | dict[str, Any],
        *,
        title: str = "More information needed",
        reason: str = "",
        ui_schema: dict[str, Any] | None = None,
        timeout: float = 300.0,
    ) -> BaseModel | dict[str, Any]:
        """Pause the skill until the caller submits structured input.

        ``schema`` may be a Pydantic model class or a JSON Schema object.
        The runtime emits ``kind="input_request"``; hosts render the schema
        as a form and POST the submitted JSON back. When a Pydantic model is
        supplied, the response is validated and returned as that model.
        """
        if isinstance(schema, type) and issubclass(schema, BaseModel):
            json_schema = schema.model_json_schema()
            raw = await _InputRegistry.collect(
                self,
                title=title,
                reason=reason,
                schema=json_schema,
                ui_schema=ui_schema,
                timeout=timeout,
            )
            return schema.model_validate(raw)
        if not isinstance(schema, dict):
            raise TypeError("collect schema must be a BaseModel class or JSON Schema dict")
        raw = await _InputRegistry.collect(
            self,
            title=title,
            reason=reason,
            schema=schema,
            ui_schema=ui_schema,
            timeout=timeout,
        )
        return TypeAdapter(dict[str, Any]).validate_python(raw)

    @classmethod
    def answer(cls, question_id: str, answer: str) -> bool:
        """Resolve a pending :meth:`ask` from outside (e.g. an HTTP handler)."""
        return _AskRegistry.answer(question_id, answer)

    @classmethod
    def submit_input(cls, request_id: str, value: dict[str, Any]) -> bool:
        """Resolve a pending :meth:`collect` from an HTTP/runtime handler."""
        return _InputRegistry.submit(request_id, value)

    async def request_scope(
        self,
        *,
        reason: str,
        read: Sequence[str] = (),
        write_prefix: str | None = None,
        ttl_seconds: int = 60,
        mode: str = "read_only",
        timeout: float = 60.0,
    ) -> Grant:
        """Ask the platform for a scope expansion mid-skill.

        Pauses execution until the platform replies. On approve, the new
        :class:`Grant` is verified, installed on ``ctx.workspace``, and
        returned. On deny, raises :class:`ScopeDenied`.

        The skill must opt in by declaring ``allow_scope_expansion=True`` on
        its ``@skill`` decorator — otherwise raises
        :class:`ScopeExpansionNotAllowed`.
        """
        if not getattr(self, "_scope_expansion_allowed", False):
            raise ScopeExpansionNotAllowed(
                "this skill did not opt in to scope expansion; "
                "set @skill(allow_scope_expansion=True) to enable"
            )
        if not reason.strip():
            raise ValueError("request_scope: reason is required")
        grant = await _ScopeRegistry.request(
            self,
            reason=reason,
            read=read,
            write_prefix=write_prefix,
            ttl_seconds=ttl_seconds,
            mode=mode,
            timeout=timeout,
        )
        # Install the grant on whatever workspace this context has, if any.
        ws = getattr(self, "_workspace", None)
        if ws is not None and hasattr(ws, "install_grant"):
            try:
                ws.install_grant(grant)
            except NotImplementedError:
                # Workspace doesn't support live grant install; the skill is
                # still free to use the returned grant value directly.
                pass
        return grant

    @classmethod
    def resolve_scope_grant(cls, request_id: str, grant_token: str) -> bool:
        """Resolve a pending :meth:`request_scope` with a fresh signed grant.

        Called by the runtime adapter when the platform POSTs the new grant
        back to the agent (see ``serve/asgi.py``'s ``/scope-grants/{id}``).
        """
        return _ScopeRegistry.resolve(request_id, grant_token)

    @classmethod
    def deny_scope(cls, request_id: str, reason: str) -> bool:
        """Resolve a pending :meth:`request_scope` with a deny reason."""
        return _ScopeRegistry.deny(request_id, reason)

    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,
        sandbox: SandboxClient | None = None,
        a2a: A2AClient | None = None,
        discover: DiscoveryClient | None = None,
        on_event: Any | None = None,  # async (AgentEvent) -> None
    ) -> None:
        self.task_id = task_id
        self.auth = auth
        self._secrets: dict[str, str] = dict(secrets or {})
        self._workspace = workspace
        self._sandbox = sandbox
        self._a2a = a2a
        self._discover = discover
        self._cancel = asyncio.Event()
        self.events: list[AgentEvent] = []
        self.artifacts: dict[str, bytes] = {}
        self._on_event = on_event

    @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

    @property
    def sandbox(self) -> SandboxClient:
        if self._sandbox is None:
            raise SandboxUnavailable(
                "no sandbox client attached to this context; "
                "the runtime layer must provision one"
            )
        return self._sandbox

    @property
    def discover(self) -> DiscoveryClient:
        if self._discover is None:
            raise PermissionError(
                "no discovery client attached; runtime must provision one"
            )
        return self._discover

    def _a2a_client(self) -> A2AClient:
        if self._a2a is None:
            raise PermissionError(
                "no A2A client attached; runtime must provision one before "
                "ctx.call(...) can be used"
            )
        return self._a2a

    async def emit_event(self, event: AgentEvent) -> None:
        self.events.append(event)
        if self._on_event is not None:
            await self._on_event(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
