"""Jobs operation helpers + operation handles.

This mirrors the reasoners layout:
- gateways return data-only records (JobRecord / Operation)
- clients expose high-level lifecycle methods
- operation handles provide polling helpers (wait/result) around job lifecycle state
"""

from __future__ import annotations

from dataclasses import dataclass
from types import TracebackType
from typing import Any, Protocol

from .constants import JOB_TERMINAL_STATES
from .errors import JobNotReadyError
from .models import Operation, OperationStatus
from .models import Job
from .util import JobIncludeDict


class _JobsClientAsyncProto(Protocol):
    async def get_operation(self, reasoner_type: str, job_id: str) -> Operation: ...

    async def wait(
        self,
        reasoner_type: str,
        job_id: str,
        *,
        timeout_s: float = 900,
        poll_interval_s: float = 0.5,
        include: JobIncludeDict | None = None,
    ) -> Job: ...

    async def get(self, reasoner_type: str, job_id: str, *, include: JobIncludeDict | None = None) -> Job | None: ...


class _JobsClientSyncProto(Protocol):
    def get_operation(self, reasoner_type: str, job_id: str) -> Operation: ...

    def wait(
        self,
        reasoner_type: str,
        job_id: str,
        *,
        timeout_s: float = 900,
        poll_interval_s: float = 0.5,
        include: JobIncludeDict | None = None,
    ) -> Job: ...

    def get(self, reasoner_type: str, job_id: str, *, include: JobIncludeDict | None = None) -> Job | None: ...


def _operation_from_job(op_id: str, job: Job | None) -> Operation:
    if job is None:
        return Operation(id=op_id, status=OperationStatus.PENDING, raw={"note": "job not found yet"})

    state = (job.state or "").upper()
    raw = {"state": state, "reasoner_type": job.reasoner_type, "name": job.name}

    if state in JOB_TERMINAL_STATES:
        if state == "COMPLETED":
            return Operation(id=op_id, status=OperationStatus.SUCCEEDED, raw=raw)
        if state == "CANCELED":
            return Operation(id=op_id, status=OperationStatus.CANCELLED, raw=raw)
        # FAILED / ABORTED (etc)
        err = job.abort_reason or f"job ended in state={state}"
        return Operation(id=op_id, status=OperationStatus.FAILED, error=err, raw=raw)

    return Operation(id=op_id, status=OperationStatus.RUNNING, raw=raw)


def operation_from_job(op_id: str, job: Job | None) -> Operation:
    """Public helper to derive an `Operation` from a `Job` view."""
    return _operation_from_job(op_id, job)


@dataclass
class JobOperation:
    """Async operation handle returned by `JobsClient.create(...)`."""

    client: _JobsClientAsyncProto
    id: str
    metadata: dict[str, Any]
    _result: Job | None = None

    def _reasoner_type(self) -> str:
        t = self.metadata.get("reasoner_type")
        if not isinstance(t, str) or not t.strip():
            raise JobNotReadyError("JobOperation metadata is missing required key: 'reasoner_type'")
        return t

    async def status(self) -> Operation:
        return await self.client.get_operation(self._reasoner_type(), self.id)

    async def is_done(self) -> bool:
        return (await self.status()).done()

    async def wait(
        self,
        *,
        timeout_s: float = 900,
        poll_interval_s: float = 0.5,
        include: JobIncludeDict | None = None,
    ) -> Job:
        if self._result is not None:
            return self._result
        j = await self.client.wait(
            self._reasoner_type(),
            self.id,
            timeout_s=timeout_s,
            poll_interval_s=poll_interval_s,
            include=include,
        )
        self._result = j
        return j

    async def result(self, *, include: JobIncludeDict | None = None) -> Job:
        if self._result is not None:
            return self._result
        op = await self.status()
        if not op.done():
            raise JobNotReadyError(f"Operation not complete yet: id={self.id} status={op.status}")
        job = await self.client.get(self._reasoner_type(), self.id, include=include)
        if job is None:
            raise JobNotReadyError(f"Operation complete but job not found: id={self.id}")
        self._result = job
        return job

    async def __aenter__(self) -> Job:
        return await self.wait()

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        return None

    def __await__(self):
        return self.wait().__await__()

    def __repr__(self) -> str:
        return f"JobOperation(id={self.id!r}, metadata={self.metadata!r})"


@dataclass
class JobOperationSync:
    """Sync operation handle returned by `JobsClientSync.create(...)`."""

    client: _JobsClientSyncProto
    id: str
    metadata: dict[str, Any]
    _result: Job | None = None

    def _reasoner_type(self) -> str:
        t = self.metadata.get("reasoner_type")
        if not isinstance(t, str) or not t.strip():
            raise JobNotReadyError("JobOperationSync metadata is missing required key: 'reasoner_type'")
        return t

    def status(self) -> Operation:
        return self.client.get_operation(self._reasoner_type(), self.id)

    def is_done(self) -> bool:
        return self.status().done()

    def wait(
        self,
        *,
        timeout_s: float = 900,
        poll_interval_s: float = 0.5,
        include: JobIncludeDict | None = None,
    ) -> Job:
        if self._result is not None:
            return self._result
        j = self.client.wait(
            self._reasoner_type(),
            self.id,
            timeout_s=timeout_s,
            poll_interval_s=poll_interval_s,
            include=include,
        )
        self._result = j
        return j

    def result(self, *, include: JobIncludeDict | None = None) -> Job:
        if self._result is not None:
            return self._result
        op = self.status()
        if not op.done():
            raise JobNotReadyError(f"Operation not complete yet: id={self.id} status={op.status}")
        job = self.client.get(self._reasoner_type(), self.id, include=include)
        if job is None:
            raise JobNotReadyError(f"Operation complete but job not found: id={self.id}")
        self._result = job
        return job

    def __repr__(self) -> str:
        return f"JobOperationSync(id={self.id!r}, metadata={self.metadata!r})"


__all__ = ["JobOperation", "JobOperationSync", "operation_from_job"]

