"""Typed protocol objects for the Sandai operator client."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, TypedDict

from celery.result import AsyncResult

if TYPE_CHECKING:
    from .async_transport import AsyncTransportAdapter


class InputArtifactDict(TypedDict, total=False):
    """Input artifact payload accepted by the client."""

    uri: str
    slot: str
    meta: Dict[str, Any]


class TaskMetaDict(TypedDict, total=False):
    """Task metadata sent to the operator."""

    persistent_output: bool


class TaskPayloadDict(TypedDict, total=False):
    """Normalized task payload sent through Celery."""

    meta: Dict[str, Any]
    artifacts: List[Dict[str, Any]]
    options: Dict[str, Any]
    deadline: float
    cancel_key: str
    task_type: str


class ErrorInfoDict(TypedDict, total=False):
    """Operator error payload."""

    error_type: str
    error_message: str
    error_details: Dict[str, Any]


@dataclass(frozen=True)
class OperatorTaskResult:
    """Normalized successful task result."""

    artifacts: List[Dict[str, Any]]
    results: Dict[str, Any]

    def to_legacy_tuple(self) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
        return self.artifacts, self.results


@dataclass(frozen=True)
class OperatorDescription:
    """Normalized operator description payload."""

    doc: Any
    options_schema: Any
    results_schema: Any

    def to_dict(self) -> Dict[str, Any]:
        return {
            "doc": self.doc,
            "options_schema": self.options_schema,
            "results_schema": self.results_schema,
        }


@dataclass(frozen=True)
class SubmittedTask:
    """Structured batch submission result."""

    index: int
    async_result: Optional[AsyncResult] = None
    error: Optional[Exception] = None

    @property
    def submitted(self) -> bool:
        return self.async_result is not None and self.error is None

    @property
    def task_id(self) -> Optional[str]:
        if self.async_result is None:
            return None
        return getattr(self.async_result, "id", None)


@dataclass(frozen=True)
class AsyncTaskHandle:
    """Async task handle backed by a transport adapter."""

    async_result: Any
    transport: "AsyncTransportAdapter"

    @property
    def id(self) -> Optional[str]:
        return getattr(self.async_result, "id", None)

    @property
    def task_id(self) -> Optional[str]:
        return self.id

    async def get(self, timeout: Optional[int] = None) -> Any:
        return await self.transport.get_result(self.async_result, timeout=timeout or 0)

    async def forget(self) -> None:
        await self.transport.forget(self.async_result)

    async def ready(self) -> bool:
        return await self.transport.result_exists(self.async_result)

    async def state(self) -> str:
        return await self.transport.get_state(self.async_result)


@dataclass(frozen=True)
class AsyncSubmittedTask:
    """Structured async batch submission result."""

    index: int
    handle: Optional[AsyncTaskHandle] = None
    error: Optional[Exception] = None

    @property
    def submitted(self) -> bool:
        return self.handle is not None and self.error is None

    @property
    def task_id(self) -> Optional[str]:
        if self.handle is None:
            return None
        return self.handle.task_id


@dataclass(frozen=True)
class BatchTaskResult:
    """Structured batch execution result."""

    index: int
    task_id: Optional[str]
    submitted: bool
    artifacts: List[Dict[str, Any]] = field(default_factory=list)
    results: Dict[str, Any] = field(default_factory=dict)
    error: Optional[Exception] = None

    @property
    def success(self) -> bool:
        return self.submitted and self.error is None

    def to_legacy_tuple(self) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
        if self.success:
            return self.artifacts, self.results
        return [], {}


@dataclass(frozen=True)
class HealthCheckResult:
    """Detailed transport/protocol probe result."""

    submit_ok: bool
    retrieval_ok: bool
    response_ok: bool
    task_id: Optional[str] = None
    error: Optional[str] = None
