"""
Execution plan data model — the contract between cloud and runtime.

This module defines the two core types that cross the trust boundary:

    PlanStep     — A single unit of work (call a model, call an agent).
    ExecutionPlan — An ordered sequence of steps with auth, TTL, and HMAC.

Design decisions:
    - Dataclasses over Pydantic: zero dependencies, works on vanilla Python 3.10+.
      The runtime is a thin client — we don't impose heavy deps.
    - Frozen steps: once the cloud emits a plan, nothing on the runtime side
      should mutate it. Steps are frozen; the plan itself is not (the executor
      needs to attach the verified signature after validation).
    - Explicit step types: "local", "cloud", "local_agent", "remote_agent".
      The runtime uses the type to pick the right bridge (Local, Cloud, Agent).
    - context_refs: symbolic references to earlier outputs. The executor resolves
      "step_1_output" at runtime — the plan never contains actual data.
"""

from __future__ import annotations

import time
import uuid
from dataclasses import dataclass, field
from typing import Any

# ── Protocol ────────────────────────────────────────────────────────────

PROTOCOL_VERSION: int = 1

VALID_STEP_TYPES: frozenset[str] = frozenset({
    "local",          # Ollama / local model on user's machine
    "cloud",          # Claude / Gemini / GPT via user's BYOK key
    "local_agent",    # Open-source agent running locally in sandbox
    "remote_agent",   # Commercial agent hosted by creator
})

VALID_AUTH_TYPES: frozenset[str] = frozenset({
    "none",               # Local models — no auth needed
    "byok",               # User's own API key (Claude, Gemini, GPT)
    "marketplace_token",  # Brynq-issued token for commercial agents
})

DEFAULT_TTL_SECONDS: int = 300  # 5 minutes


# ── PlanStep ────────────────────────────────────────────────────────────

@dataclass(frozen=True)
class PlanStep:
    """A single executable step in an execution plan.

    Attributes:
        step_id:      Unique within the plan ("step_1", "step_2", ...).
        type:         Execution target category. Must be one of VALID_STEP_TYPES.
        backend:      Engine identifier — "ollama", "claude", "gemini", "gpt", "agent".
        model:        Specific model or agent version string.
        prompt:       The final prompt text (optimized by the cloud planner).
        context_refs: Symbolic references to prior step outputs or "original_input".
        config:       Passthrough config — temperature, max_tokens, system_prompt, etc.
        agent_id:     Marketplace agent ID (agents only).
        endpoint:     Target URL — "localhost:9100" or "https://api.acme.com/v1" (agents only).
        auth_type:    How to authenticate — "none", "byok", "marketplace_token".
    """

    step_id: str
    type: str
    backend: str
    model: str
    prompt: str
    context_refs: list[str] = field(default_factory=list)
    config: dict[str, Any] = field(default_factory=dict)
    # Agent-specific (None for plain model steps)
    agent_id: str | None = None
    endpoint: str | None = None
    auth_type: str | None = None

    def __post_init__(self) -> None:
        if self.type not in VALID_STEP_TYPES:
            raise ValueError(
                f"Invalid step type {self.type!r}. "
                f"Must be one of: {', '.join(sorted(VALID_STEP_TYPES))}"
            )
        if self.auth_type is not None and self.auth_type not in VALID_AUTH_TYPES:
            raise ValueError(
                f"Invalid auth_type {self.auth_type!r}. "
                f"Must be one of: {', '.join(sorted(VALID_AUTH_TYPES))}"
            )
        if self.type in ("local_agent", "remote_agent"):
            if not self.agent_id:
                raise ValueError(f"Step type {self.type!r} requires agent_id")
            if not self.endpoint:
                raise ValueError(f"Step type {self.type!r} requires endpoint")
        if not self.step_id:
            raise ValueError("step_id cannot be empty")
        if not self.prompt:
            raise ValueError("prompt cannot be empty")


# ── ExecutionPlan ───────────────────────────────────────────────────────

@dataclass
class ExecutionPlan:
    """A signed, time-bounded sequence of steps for the runtime to execute.

    Created by the cloud planner, serialized as JSON, sent over HTTPS,
    verified by the runtime before execution.

    Attributes:
        plan_id:        Globally unique identifier (UUID4).
        user_id:        The authenticated user this plan belongs to.
        steps:          Ordered list of PlanSteps.
        original_input: The user's original request (for context_ref resolution).
        created_at:     Unix timestamp of plan creation.
        expires_at:     Unix timestamp after which the plan is invalid.
        signature:      HMAC-SHA256 hex digest (set by signer, verified by runtime).
        version:        Protocol version. Runtime rejects unknown versions.
    """

    plan_id: str
    user_id: str
    steps: list[PlanStep]
    original_input: str
    created_at: float
    expires_at: float
    signed_at: str = ""
    signature: str = ""
    version: int = PROTOCOL_VERSION

    def __post_init__(self) -> None:
        if self.version != PROTOCOL_VERSION:
            raise ValueError(
                f"Unsupported protocol version {self.version}. "
                f"Expected {PROTOCOL_VERSION}."
            )
        if not self.steps:
            raise ValueError("ExecutionPlan must contain at least one step")

    @staticmethod
    def new(
        user_id: str,
        steps: list[PlanStep],
        original_input: str,
        ttl_seconds: int = DEFAULT_TTL_SECONDS,
    ) -> ExecutionPlan:
        """Factory: create a fresh plan with auto-generated ID and timestamps.

        Raises:
            ValueError: If ttl_seconds is not positive.
        """
        if ttl_seconds <= 0:
            raise ValueError("ttl_seconds must be positive")
        now = time.time()
        return ExecutionPlan(
            plan_id=str(uuid.uuid4()),
            user_id=user_id,
            steps=steps,
            original_input=original_input,
            created_at=now,
            expires_at=now + ttl_seconds,
        )
