"""Shared contracts and job-local accounting for Agent-mode workers."""

import threading
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Protocol


@dataclass(frozen=True)
class AgentScope:
    """A bounded work package passed to exactly one specialist agent."""

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


@dataclass
class AgentResult:
    """Artifacts returned by a specialist for collection by the orchestrator."""

    scope_id: str
    artifacts: List[Dict[str, Any]] = field(default_factory=list)
    error: str = ""


@dataclass
class AgentUsage:
    """Thread-safe usage accounting owned by one Agent pipeline run."""

    usd: float = 0.0
    calls: int = 0
    cached: int = 0
    by_stage: Dict[str, Dict[str, Any]] = field(default_factory=dict)
    models: Dict[str, set] = field(default_factory=lambda: {
        "vision": set(),
        "text_cloud": set(),
    })
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)

    def record(
        self,
        stage: str,
        model: str,
        usd: float = 0.0,
        cached: bool = False,
        has_images: bool = False,
    ) -> None:
        with self._lock:
            bucket = self.by_stage.setdefault(
                stage, {"usd": 0.0, "calls": 0, "cached": 0}
            )
            if cached:
                self.cached += 1
                bucket["cached"] += 1
            else:
                self.calls += 1
                self.usd += usd
                bucket["calls"] += 1
                bucket["usd"] += usd
            family = "vision" if has_images else "text_cloud"
            self.models[family].add(model)

    def snapshot(self) -> Dict[str, Any]:
        with self._lock:
            return {
                "usd": self.usd,
                "calls": self.calls,
                "cached": self.cached,
                "by_stage": {k: dict(v) for k, v in self.by_stage.items()},
                "models": {
                    "vision": sorted(self.models["vision"]),
                    "text_local": [],
                    "text_cloud": sorted(self.models["text_cloud"]),
                    "fallback_count": 0,
                },
            }


class ScopedAgent(Protocol):
    """Protocol implemented by each future specialist agent."""

    name: str

    def run(self, scope: AgentScope) -> AgentResult:
        ...


def failure(scope: AgentScope, error: Exception) -> AgentResult:
    """Convert a worker exception into a non-fatal scoped result."""
    return AgentResult(scope_id=scope.scope_id, error=str(error))