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


TemplateUpdatePolicy = Literal["none", "notify", "propose", "auto_patch"]


class TemplateLineage(BaseModel):
    """Opt-in source-template lineage and update policy.

    Agents generated from templates can publish this on their Agent Card so
    platforms can detect available template updates and decide how to handle
    them. The policy is intentionally advisory: callers/platforms remain the
    authority for whether an update is proposed, applied, reviewed, or denied.
    """

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

    schema_version: str = "2026-06-02"
    template_ref: str | None = None
    template_version: str | None = None
    template_digest: str | None = None
    source_agent: str | None = None
    source_agent_version: str | None = None
    source_repo_url: str | None = None
    source_revision: str | None = None
    instance_id: str | None = None
    instance_version: str | None = None
    update_policy: TemplateUpdatePolicy = "none"
    update_channel: str | None = None
    migration_skill: str | None = None

    @classmethod
    def none(cls) -> "TemplateLineage":
        return cls()

    @property
    def enabled(self) -> bool:
        return any(
            (
                self.template_ref,
                self.template_version,
                self.template_digest,
                self.source_agent,
                self.source_agent_version,
                self.source_repo_url,
                self.source_revision,
                self.instance_id,
                self.instance_version,
                self.update_policy != "none",
                self.update_channel,
                self.migration_skill,
            )
        )

    @classmethod
    def from_mapping(cls, data: Mapping[str, Any] | str | None) -> "TemplateLineage":
        if data is None:
            return cls.none()
        if isinstance(data, str):
            return cls(template_ref=data)
        if isinstance(data, Mapping):
            return cls.model_validate(dict(data))
        raise ValueError("template_lineage must be a mapping or template ref string")

    @field_validator(
        "template_ref",
        "template_version",
        "template_digest",
        "source_agent",
        "source_agent_version",
        "source_repo_url",
        "source_revision",
        "instance_id",
        "instance_version",
        "update_channel",
        "migration_skill",
    )
    @classmethod
    def _clean_optional_text(cls, value: str | None) -> str | None:
        cleaned = str(value or "").strip()
        return cleaned or None

    @field_validator("schema_version")
    @classmethod
    def _clean_schema_version(cls, value: str) -> str:
        return str(value or "").strip() or "2026-06-02"

    @field_validator("update_policy", mode="before")
    @classmethod
    def _clean_update_policy(cls, value: Any) -> str:
        cleaned = str(value or "none").strip().lower().replace("-", "_")
        return cleaned or "none"

    @model_validator(mode="after")
    def _require_source_for_updates(self) -> "TemplateLineage":
        if self.update_policy != "none" and not (self.template_ref or self.source_agent):
            raise ValueError(
                "template lineage update_policy requires template_ref or source_agent"
            )
        return self

    @model_serializer(mode="wrap")
    def _dump_without_empty_fields(self, handler: Any) -> dict[str, Any]:
        data = handler(self)
        return {key: value for key, value in data.items() if value is not None}


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 parsed ``a2a.yaml`` platform metadata to an agent class.

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

    raw = dict(config or {})
    manifest = MetaAgentManifest.from_mapping(config)
    if manifest.enabled:
        setattr(agent_cls, "meta_agent_manifest", manifest)
    lineage = TemplateLineage.from_mapping(
        raw.get("template_lineage") or raw.get("template")
    )
    if lineage.enabled:
        setattr(agent_cls, "template_lineage", lineage)
    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."""
