"""Declarative runtime/deployment metadata.

These types describe *how* the platform should run an agent: lifecycle,
state needs, isolation level, resource budget, egress policy. They are
read by the deployer and by the registry; agent code itself should not
depend on which runtime is selected.
"""
from __future__ import annotations

from enum import Enum
from typing import Any, Literal, Mapping

from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    NonNegativeInt,
    PositiveInt,
    field_validator,
    model_serializer,
    model_validator,
)


class Lifecycle(str, Enum):
    """How long an instance of the agent process lives."""

    EPHEMERAL = "ephemeral"  # spawned per-invocation, torn down after
    SESSION = "session"  # spawned per-session, kept until the session ends
    WARM = "warm"  # long-running service, multiplexed across callers


class State(str, Enum):
    """What kind of state the agent retains between invocations."""

    NONE = "none"  # purely functional
    SESSION = "session"  # in-memory state for the lifetime of a session
    DURABLE = "durable"  # persisted across restarts (storage required)


class Sandbox(str, Enum):
    """Isolation level. The platform always runs agents under microsandbox.

    Modeled as an enum (rather than a constant) so the wire format stays
    stable if more isolation tiers are added later, but only one value is
    currently valid: every agent runs in a microvm-class sandbox.
    """

    MICROSANDBOX = "microsandbox"


class Resources(BaseModel):
    """Resource budget hint for the deployer."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    cpu: str = "100m"  # k8s-style CPU spec, e.g. "500m", "2"
    memory: str = "256Mi"  # k8s-style memory, e.g. "512Mi", "4Gi"
    gpu: NonNegativeInt = 0
    max_runtime_seconds: PositiveInt = 600


class SkillPolicy(BaseModel):
    """Per-skill operational policy advertised on the Agent Card."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    timeout_seconds: float | None = None
    idempotent: bool = False
    max_retries: NonNegativeInt = 0
    cost_class: str | None = None  # informational, e.g. "cheap" / "expensive"
    # When True, the skill is allowed to call ``ctx.request_scope(...)`` to
    # negotiate additional workspace access mid-execution. Off by default —
    # platform refuses scope negotiation otherwise.
    allow_scope_expansion: bool = False

    # Optional first-grant hints. The caller remains the authority: it may
    # clamp, narrow, deny, or ignore these values before minting a grant.
    grant_mode: str | None = None
    grant_allow_patterns: tuple[str, ...] = ()
    grant_deny_patterns: tuple[str, ...] = ()
    grant_outputs_prefix: str | None = None
    grant_write_prefixes: tuple[str, ...] = ()
    grant_ttl_seconds: PositiveInt | None = None
    grant_run_timeout_seconds: PositiveInt | None = None
    grant_approval_timeout_seconds: PositiveInt | None = None
    grant_scope_approval_timeout_seconds: PositiveInt | None = None


class EgressPolicy(BaseModel):
    """What external hosts the agent is allowed to talk to."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    allow_hosts: tuple[str, ...] = ()
    allow_internal_services: tuple[str, ...] = ()  # e.g. cluster service DNS
    deny_internet_by_default: bool = True


class LLMProvisioning(str, Enum):
    """Where the agent's LLM credentials come from.

    Drives both runtime behaviour (``ctx.llm`` resolution) and what the
    marketplace can charge the caller.
    """

    PLATFORM = "platform"
    """Use the platform's shared LiteLLM gateway (free during beta;
    folded into platform pricing in prod). Authors don't need to know
    or care about keys."""

    PLATFORM_OR_CALLER_PROVIDED = "platform_or_caller_provided"
    """Use caller-provided credentials when the caller selected them;
    otherwise fall back to a scoped platform LiteLLM grant. Intended for
    trusted platform agents that should work in both billing modes."""

    CALLER_PROVIDED = "caller_provided"
    """Use the caller's own LLM credentials, forwarded by the CP in the
    invoke body. The author's per-call price covers the skill IP /
    compute only; the LLM bill goes to the caller's provider directly."""

    AGENT_BYOK = "agent_byok"
    """Author brings their own LLM key (set via env / Secret). The
    per-call price typically includes a markup on the LLM cost."""


class Pricing(BaseModel):
    """Per-agent pricing declaration. Surfaced on the Card; the
    marketplace uses it to display "this much per call" to potential
    callers.
    """

    model_config = ConfigDict(extra="forbid", frozen=True)

    price_per_call_usd: float = 0.0
    """Flat charge per skill invocation, in USD. ``0.0`` = free."""

    caller_pays_llm: bool = False
    """If True, the caller's LLM provider eats the inference cost
    directly (matches ``LLMProvisioning.CALLER_PROVIDED``). If False the
    author's price is expected to cover LLM cost too."""

    notes: str = ""
    """Free-form marketing text the marketplace will render next to the
    price (e.g. "first 100 calls free")."""

    compute: dict[str, float] | None = None
    """Platform-populated compute quote, derived from runtime resources.

    Agent authors should not set this directly; the control plane fills it
    into registry/discovery responses so SDK clients can validate enriched
    Agent Cards.
    """

    total_usd: float | None = None
    """Platform-populated total per-call quote, including compute and markup."""

    @model_serializer(mode="wrap")
    def _dump_without_empty_quote_fields(self, handler: Any) -> dict[str, Any]:
        data = handler(self)
        if data.get("compute") is None:
            data.pop("compute", None)
        if data.get("total_usd") is None:
            data.pop("total_usd", None)
        return data


class CompositionSubAgent(BaseModel):
    """One callable dependency declared by a composable/meta agent."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    name: str | None = None
    tag: str | None = None
    version: str | None = None
    skills: tuple[str, ...] = ()
    default_args: dict[str, Any] = Field(default_factory=dict)
    required: bool = True

    @field_validator("name", "tag", "version")
    @classmethod
    def _clean_optional_text(cls, value: str | None) -> str | None:
        cleaned = str(value or "").strip()
        return cleaned or None

    @field_validator("skills", mode="before")
    @classmethod
    def _clean_skills(cls, value: Any) -> tuple[str, ...]:
        if value is None:
            return ()
        if isinstance(value, str):
            raw_items = [value]
        elif isinstance(value, (list, tuple, set)):
            raw_items = list(value)
        else:
            raise ValueError("skills must be a string or list of strings")
        seen: set[str] = set()
        out: list[str] = []
        for item in raw_items:
            skill = str(item or "").strip()
            if not skill or skill in seen:
                continue
            seen.add(skill)
            out.append(skill)
        return tuple(out)

    @model_validator(mode="after")
    def _require_binding(self) -> "CompositionSubAgent":
        if not self.name and not self.tag:
            raise ValueError("composition sub-agent requires name or tag")
        return self

    def public_payload(self) -> dict[str, Any]:
        payload: dict[str, Any] = {"required": self.required}
        if self.name:
            payload["name"] = self.name
        if self.tag:
            payload["tag"] = self.tag
        if self.version:
            payload["version"] = self.version
        if self.skills:
            payload["skills"] = list(self.skills)
        if self.default_args:
            payload["default_arg_keys"] = sorted(str(key) for key in self.default_args)
        return payload


class AgentComposition(BaseModel):
    """Declarative composition block from ``a2a.yaml``."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    sub_agents: tuple[CompositionSubAgent, ...] = ()
    planning: Literal["llm_dag", "deterministic_dag"] = "llm_dag"
    max_nodes: PositiveInt = 8
    max_parallel: PositiveInt = 3
    max_replans: NonNegativeInt = 1

    @model_validator(mode="before")
    @classmethod
    def _aliases(cls, value: Any) -> Any:
        if isinstance(value, (list, tuple)):
            return {"sub_agents": value}
        if isinstance(value, dict):
            raw = dict(value)
            if "sub_agents" not in raw:
                for alias in ("agents", "children"):
                    if alias in raw:
                        raw["sub_agents"] = raw.pop(alias)
                        break
            return raw
        return value

    def public_payload(self) -> dict[str, Any]:
        return {
            "planning": self.planning,
            "max_nodes": self.max_nodes,
            "max_parallel": self.max_parallel,
            "max_replans": self.max_replans,
            "sub_agents": [item.public_payload() for item in self.sub_agents],
        }


class AgentGoal(BaseModel):
    """Durable objective advertised by a composable/meta agent."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    objective: str = ""
    success_criteria: tuple[str, ...] = ()
    constraints: tuple[str, ...] = ()

    @field_validator("objective")
    @classmethod
    def _clean_objective(cls, value: str) -> str:
        return str(value or "").strip()

    @field_validator("success_criteria", "constraints", mode="before")
    @classmethod
    def _clean_string_tuple(cls, value: Any) -> tuple[str, ...]:
        if value is None:
            return ()
        if isinstance(value, str):
            raw_items = [value]
        elif isinstance(value, (list, tuple, set)):
            raw_items = list(value)
        else:
            raise ValueError("must be a string or list of strings")
        return tuple(str(item).strip() for item in raw_items if str(item or "").strip())

    def public_payload(self) -> dict[str, Any]:
        payload: dict[str, Any] = {}
        if self.objective:
            payload["objective"] = self.objective
        if self.success_criteria:
            payload["success_criteria"] = list(self.success_criteria)
        if self.constraints:
            payload["constraints"] = list(self.constraints)
        return payload


MemoryTier = Literal["files", "kv", "vector"]


class AgentMemory(BaseModel):
    """Long-term memory declaration for composable/meta agents."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    tiers: tuple[MemoryTier, ...] = ()
    namespace: str | None = None
    scope: Literal["agent", "user", "thread"] = "agent"
    retention: Literal["ephemeral", "durable"] = "durable"

    @field_validator("tiers", mode="before")
    @classmethod
    def _clean_tiers(cls, value: Any) -> tuple[str, ...]:
        if value is None:
            return ()
        if isinstance(value, str):
            raw_items = [value]
        elif isinstance(value, (list, tuple, set)):
            raw_items = list(value)
        else:
            raise ValueError("tiers must be a string or list of strings")
        seen: set[str] = set()
        out: list[str] = []
        for item in raw_items:
            tier = str(item or "").strip().lower()
            if not tier or tier in seen:
                continue
            seen.add(tier)
            out.append(tier)
        return tuple(out)

    @field_validator("namespace")
    @classmethod
    def _clean_namespace(cls, value: str | None) -> str | None:
        cleaned = str(value or "").strip()
        return cleaned or None

    def public_payload(self) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "tiers": list(self.tiers),
            "scope": self.scope,
            "retention": self.retention,
        }
        if self.namespace:
            payload["namespace"] = self.namespace
        return payload


class MetaAgentManifest(BaseModel):
    """Typed meta-agent contract parsed from ``a2a.yaml`` or class attrs."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    composition: AgentComposition | None = None
    goal: AgentGoal | None = None
    memory: AgentMemory | None = None

    @property
    def enabled(self) -> bool:
        return bool(self.composition or self.goal or self.memory)

    @classmethod
    def from_mapping(cls, data: Mapping[str, Any] | None) -> "MetaAgentManifest":
        raw = dict(data or {})
        composition = raw.get("composition")
        goal = raw.get("goal")
        memory = raw.get("memory")
        return cls(
            composition=AgentComposition.model_validate(composition)
            if isinstance(composition, (dict, list, tuple))
            else None,
            goal=AgentGoal.model_validate({"objective": goal})
            if isinstance(goal, str)
            else AgentGoal.model_validate(goal)
            if isinstance(goal, dict)
            else None,
            memory=AgentMemory.model_validate({"tiers": memory})
            if isinstance(memory, (str, list, tuple))
            else AgentMemory.model_validate(memory)
            if isinstance(memory, dict)
            else None,
        )

    def public_payload(self) -> dict[str, Any]:
        payload: dict[str, Any] = {"schema_version": "2026-06-02"}
        if self.composition is not None:
            payload["composition"] = self.composition.public_payload()
        if self.goal is not None:
            goal = self.goal.public_payload()
            if goal:
                payload["goal"] = goal
        if self.memory is not None:
            payload["memory"] = self.memory.public_payload()
        return payload if len(payload) > 1 else {}


def apply_project_manifest(
    agent_cls: type[Any],
    config: Mapping[str, Any] | None,
) -> MetaAgentManifest:
    """Attach a parsed ``a2a.yaml`` meta manifest to an agent class.

    ``a2a run`` receives only an entrypoint, but the working directory still
    contains ``a2a.yaml``. Applying the parsed manifest here makes local and
    deployed Agent Cards reflect declarative composition without importing
    control-plane code.
    """

    manifest = MetaAgentManifest.from_mapping(config)
    if manifest.enabled:
        setattr(agent_cls, "meta_agent_manifest", manifest)
    return manifest


class AgentRuntime(BaseModel):
    """Aggregate runtime declaration; published on the Agent Card."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    lifecycle: Lifecycle = Lifecycle.EPHEMERAL
    state: State = State.NONE
    sandbox: Sandbox = Sandbox.MICROSANDBOX
    resources: Resources = Field(default_factory=Resources)
    concurrency: PositiveInt = 1
    egress: EgressPolicy = Field(default_factory=EgressPolicy)
    tools_used: tuple[str, ...] = ()
    llm_provisioning: LLMProvisioning = LLMProvisioning.PLATFORM
    pricing: Pricing = Field(default_factory=Pricing)
    wants_cp_jwt: bool = False
    """When True, the platform forwards the *caller's* control-plane JWT
    to this agent in the /invoke body so the skill can call back into
    user-scoped CP endpoints (/v1/me/files, /v1/agents/from-tarball,
    etc.) on the user's behalf. Only opt in for trusted platform agents
    — a JWT lets the holder do anything the user can do."""

    apt_packages: tuple[str, ...] = ()
    """Extra Debian packages to apt-install on top of the base image at
    build time. Use this when the agent needs system binaries the
    Python-only base doesn't ship — ffmpeg, imagemagick, poppler-utils,
    etc. Each entry must match ``[a-z0-9.+-]{2,64}`` (validated by the
    control plane); anything fancier (custom apt repos, --no-install-
    recommends overrides) belongs in a bespoke base image, not here."""
