"""Subagent protocol and interface definitions.

Defines the abstract interface that all subagent implementations must satisfy,
enforcing context isolation, token budgets, and structured outputs.
"""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Protocol, runtime_checkable

from cyberbrain26.orchestration.subagent.types import (
    PatternType,
    SubagentOutput,
    SubagentSpec,
)

logger = logging.getLogger(__name__)


@runtime_checkable
class SubagentProtocol(Protocol):
    """Protocol defining the interface for subagent implementations."""

    async def invoke(self, input_data: dict[str, Any]) -> SubagentOutput:
        """Execute the subagent with the given input and return structured output.

        Must respect token_budget, timeout, and return structured SubagentOutput.
        """
        ...

    def get_spec(self) -> SubagentSpec:
        """Return the specification for this subagent."""
        ...

    async def checkpoint(self) -> dict[str, Any]:
        """Return checkpoint state for recovery (Phase 5)."""
        ...

    async def recover(self, checkpoint: dict[str, Any]) -> None:
        """Restore state from a checkpoint (Phase 5)."""
        ...


class BaseSubagent(ABC):
    """Abstract base class for subagent implementations.

    Provides common functionality for LangGraph, custom, or other subagent runners.
    """

    def __init__(self, spec: SubagentSpec) -> None:
        if not isinstance(spec, SubagentSpec):
            raise TypeError(f"Expected SubagentSpec, got {type(spec)}")
        self._spec = spec
        self._state: dict[str, Any] = {}
        self._tokens_used = 0
        self._start_time = 0

    @property
    def spec(self) -> SubagentSpec:
        return self._spec

    @property
    def pattern(self) -> PatternType:
        return self._spec.pattern

    @abstractmethod
    async def invoke(self, input_data: dict[str, Any]) -> SubagentOutput:
        """Execute the subagent with the given input.

        Must be implemented by subclasses with pattern-specific logic.
        """
        raise NotImplementedError

    def get_spec(self) -> SubagentSpec:
        return self._spec

    async def checkpoint(self) -> dict[str, Any]:
        """Return checkpoint state for recovery."""
        return {
            "spec": self._spec,
            "state": self._state,
            "tokens_used": self._tokens_used,
            "start_time": self._start_time,
        }

    async def recover(self, checkpoint: dict[str, Any]) -> None:
        """Restore state from a checkpoint."""
        self._state = checkpoint.get("state", {})
        self._tokens_used = checkpoint.get("tokens_used", 0)
        self._start_time = checkpoint.get("start_time", 0)


@dataclass
class SubagentFactory:
    """Factory for creating subagents from specifications."""

    role: str
    runner_class: type[BaseSubagent]

    def create(self, spec: SubagentSpec) -> BaseSubagent:
        """Create a subagent instance with the given spec."""
        return self.runner_class(spec)


class SubagentError(Exception):
    """Base exception for subagent errors."""

    pass


class BudgetExceeded(SubagentError):
    """Raised when token budget is exceeded."""

    pass


class TimeoutExceeded(SubagentError):
    """Raised when execution exceeds timeout."""

    pass


class ContextOverflow(SubagentError):
    """Raised when context window reaches limit."""

    pass


class DelegationDenied(SubagentError):
    """Raised when delegation is denied by security gate."""

    pass