"""Reasoners operation helpers + operation handles.

This module contains backend-agnostic helpers to:
- generate stable client-side operation ids (when a backend doesn't expose one)
- derive operation status from a reasoner lifecycle state
- provide async + sync operation handle objects (`ReasonerOperation`, `ReasonerOperationSync`)

Rationale:
- Gateways return the data-only `Operation` model (what the backend reports).
- Clients wrap that into higher-level operation handles that implement polling + READY semantics.
"""

from __future__ import annotations

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

from relationalai.util.polling import poll_until

from . import constants as c
from .errors import (
    ReasonerInvalidRequestError,
    ReasonerNotReadyError,
    ReasonerTimeoutError,
)
from .models import Operation, OperationStatus, Reasoner

class _ReasonersClientAsyncProto(Protocol):
    def _polling_cfg(self) -> tuple[float, float, float]: ...

    async def get_operation(self, operation_id: str) -> Operation: ...

    async def _wait_until_ready(self, *, reasoner_name: str, reasoner_type: str, timeout_s: float) -> Reasoner: ...

    async def get(self, reasoner_type: str, reasoner_name: str) -> Reasoner | None: ...

    def _check_reasoner_ready(
        self, r: Reasoner | None, reasoner_name: str, reasoner_type: str, context: str
    ) -> Reasoner: ...

    async def _ensure_wait_finalize(self, ctx: "_EnsureContext", *, timeout_s: float) -> Reasoner: ...


class _ReasonersClientSyncProto(Protocol):
    def get_operation(self, operation_id: str) -> Operation: ...

    def wait_for_operation(self, operation_id: str, *, timeout_s: int = 900) -> Operation: ...

    def wait_until_ready(self, reasoner_type: str, reasoner_name: str, *, timeout_s: float = 900) -> Reasoner: ...

    def _raise_on_mismatch(
        self,
        *,
        reasoner_name: str,
        reasoner_type: str,
        requested: "_EnsureRequested",
        existing: Reasoner,
    ) -> None: ...


def _remaining_timeout_s(deadline: float, *, original_timeout_s: float) -> float:
    remaining = float(deadline) - time.monotonic()
    if remaining <= 0:
        raise ReasonerTimeoutError(f"Timed out after {original_timeout_s}s waiting for completion")
    return remaining


def _op_id(kind: str, *, reasoner_type: str, reasoner_name: str) -> str:
    """Stable operation id used for client-side polling when backends don't expose server-side operation ids."""

    return f"reasoner:{kind}:{reasoner_type.upper()}:{reasoner_name}"


def _parse_op_id(operation_id: str) -> tuple[str, str, str]:
    """Parse a reasoner operation id into `(kind, type, name)`.

    Args:
        operation_id: Operation id in the format `"reasoner:<kind>:<TYPE>:<NAME>"`.

    Returns:
        Tuple of `(kind, reasoner_type, reasoner_name)`.

    Raises:
        ReasonerInvalidRequestError: If the id does not match the expected format.
    """

    parts = operation_id.split(":", 3)
    if len(parts) != 4 or parts[0] != "reasoner":
        raise ReasonerInvalidRequestError(
            f"Invalid operation_id '{operation_id}'. Expected format: reasoner:<kind>:<TYPE>:<NAME>"
        )
    _, kind, rtype, rname = parts
    return kind, rtype, rname


def _op_from_reasoner(kind: str, operation_id: str, r: Reasoner | None) -> Operation:
    """Best-effort mapping of reasoner state into an operation lifecycle."""

    def _mk(
        status: str | None,
        *,
        state: str | None = None,
        error: str | None = None,
        note: str | None = None,
    ) -> Operation:
        raw: dict[str, Any] = {"kind": kind}
        if state is not None:
            raw["state"] = state
        if note is not None:
            raw["note"] = note
        return Operation(id=operation_id, status=status, error=error, raw=raw)

    if r is None:
        return _mk(c.STATUS_PENDING, note="reasoner not found yet")

    state = (r.state or "").upper()

    _SUCCEEDED: dict[str, set[str]] = {
        # Operation handles (`ReasonerOperation.wait()`) are expected to return a READY reasoner.
        # Treat SUCCEEDED as "READY now" (not merely "exists"), otherwise `wait()` can succeed
        # while the final readiness check times out.
        "create": {c.REASONER_STATE_READY},
        "ensure": {c.REASONER_STATE_READY},
        "resume": {c.REASONER_STATE_READY},
    }
    _RUNNING: dict[str, set[str]] = {
        "create": {c.REASONER_STATE_PENDING, c.REASONER_STATE_PROVISIONING},
        "ensure": {c.REASONER_STATE_PENDING, c.REASONER_STATE_PROVISIONING, c.REASONER_STATE_SUSPENDED},
        "resume": {c.REASONER_STATE_PENDING, c.REASONER_STATE_PROVISIONING},
    }
    _FAILED: dict[str, set[str]] = {
        # If a created reasoner ends up SUSPENDED, callers should not treat the create op as "ready".
        "create": {c.REASONER_STATE_FAILED, c.REASONER_STATE_ERROR, c.REASONER_STATE_SUSPENDED},
        "ensure": {c.REASONER_STATE_FAILED, c.REASONER_STATE_ERROR},
        "resume": {c.REASONER_STATE_FAILED, c.REASONER_STATE_ERROR},
    }
    _FAIL_MSG: dict[str, str] = {
        "create": "Reasoner creation failed (not READY)",
        "ensure": "Reasoner ensure failed",
        "resume": "Reasoner resume failed",
    }

    if kind in _SUCCEEDED and state in _SUCCEEDED[kind]:
        return _mk(c.STATUS_SUCCEEDED, state=state)
    if kind in _FAILED and state in _FAILED[kind]:
        return _mk(
            c.STATUS_FAILED,
            state=state,
            error=f"{_FAIL_MSG.get(kind, 'Operation failed')}: state={state}",
        )
    if kind in _RUNNING and state in _RUNNING[kind]:
        return _mk(c.STATUS_RUNNING, state=state)

    if kind in _SUCCEEDED:
        return _mk(c.STATUS_RUNNING, state=state)

    return _mk(None, state=state, note="unknown operation kind")


def _wait_for_operation_sync(fetch_op: Any, operation_id: str, *, timeout_s: int = 900) -> Operation:
    """Poll `fetch_op(operation_id)` until the operation is done or times out."""

    deadline = time.monotonic() + timeout_s
    delay = 0.05
    while True:
        op: Operation = fetch_op(operation_id)
        if op.done():
            return op
        if time.monotonic() >= deadline:
            return Operation(
                id=operation_id,
                status=c.STATUS_FAILED,
                error=f"Timed out after {timeout_s}s waiting for completion",
            )
        time.sleep(delay)
        delay = min(0.5, delay * 1.25)


async def _wait_for_operation_async(fetch_op: Any, operation_id: str, *, timeout_s: int = 900) -> Operation:
    """Poll `await fetch_op(operation_id)` until the operation is done or times out."""

    deadline = time.monotonic() + timeout_s
    delay = 0.05
    while True:
        op: Operation = await fetch_op(operation_id)
        if op.done():
            return op
        if time.monotonic() >= deadline:
            return Operation(
                id=operation_id,
                status=c.STATUS_FAILED,
                error=f"Timed out after {timeout_s}s waiting for completion",
            )
        await asyncio.sleep(delay)
        delay = min(0.5, delay * 1.25)


def _raise_on_failed_operation(op: Operation, *, operation_id: str) -> None:
    """Raise a stable error for a non-succeeded operation."""

    if (op.status or "").upper() == OperationStatus.SUCCEEDED:
        return
    err = (op.error or "").strip()
    if err.lower().startswith("timed out after"):
        raise ReasonerTimeoutError(err)
    raise ReasonerNotReadyError(
        op.error or f"Operation did not succeed: operation_id={operation_id} status={op.status}"
    )


@dataclass(frozen=True)
class _EnsureRequested:
    """Subset of requested fields that `ensure(...)` can validate on existing reasoners.

    This intentionally includes only fields exposed by current reasoner metadata
    surfaces (`get_reasoner`, `api.reasoners`, and the Direct Access equivalents).
    `await_storage_vacuum` may be surfaced through `Reasoner.settings`; when it is
    present there, `ensure(...)` can validate it against the requested value.
    """

    size: str | None
    auto_suspend_mins: int | None
    await_storage_vacuum: bool | None
    settings: dict[str, Any] | None


@dataclass(frozen=True)
class _EnsureContext:
    """Context attached to ensure operations so `await op` means READY at await-time."""

    reasoner_name: str
    reasoner_type: str
    resume: bool
    requested: _EnsureRequested


@dataclass
class ReasonerOperation:
    """Async operation handle for reasoner lifecycle actions."""

    client: _ReasonersClientAsyncProto
    id: str
    metadata: dict[str, Any]
    _result: Reasoner | None = None
    _task: asyncio.Task[Reasoner] | None = None
    _ensure_ctx: _EnsureContext | None = None

    async def status(self) -> Operation:
        """Return current operation status without blocking."""
        if self._task is not None:
            if not self._task.done():
                return Operation(id=self.id, status=OperationStatus.RUNNING)
            exc = self._task.exception()
            if exc is not None:
                return Operation(id=self.id, status=OperationStatus.FAILED, error=str(exc))
            return Operation(id=self.id, status=OperationStatus.SUCCEEDED)
        return await self.client.get_operation(self.id)

    async def is_done(self) -> bool:
        """Return True if the operation is in a terminal state."""
        if self._task is not None:
            return self._task.done()
        op = await self.status()
        return OperationStatus.done(op.status)

    async def wait(self, *, timeout_s: float = 900) -> Reasoner:
        """Wait for completion and return the resulting READY reasoner."""
        if self._result is not None:
            return self._result

        deadline = time.monotonic() + float(timeout_s)

        if self._task is not None:
            try:
                r = await asyncio.wait_for(
                    self._task, timeout=_remaining_timeout_s(deadline, original_timeout_s=float(timeout_s))
                )
            except asyncio.TimeoutError as e:
                raise ReasonerTimeoutError(f"Timed out after {timeout_s}s waiting for completion") from e
            if self._ensure_ctx is not None:
                self._result = await self.client._ensure_wait_finalize(
                    self._ensure_ctx,
                    timeout_s=_remaining_timeout_s(deadline, original_timeout_s=float(timeout_s)),
                )
                return self._result
            self._result = r
            return self._result

        initial_delay, overhead_rate, max_delay = self.client._polling_cfg()
        await poll_until(
            self.is_done,
            timeout_s=_remaining_timeout_s(deadline, original_timeout_s=float(timeout_s)),
            initial_delay=initial_delay,
            overhead_rate=overhead_rate,
            max_delay=max_delay,
            timeout_error=lambda t: ReasonerTimeoutError(f"Timed out after {t}s waiting for completion"),
        )
        return await self.result(
            timeout_s=_remaining_timeout_s(deadline, original_timeout_s=float(timeout_s))
        )

    async def result(self, *, timeout_s: float = 900) -> Reasoner:
        """Return the operation result if complete (or raise if failed/incomplete)."""
        if self._result is not None:
            return self._result

        if self._task is not None:
            if not self._task.done():
                raise ReasonerNotReadyError(f"Operation not complete yet: operation_id={self.id}")
            exc = self._task.exception()
            if exc is not None:
                raise exc.with_traceback(exc.__traceback__) from None
            if self._ensure_ctx is not None:
                self._result = await self.client._ensure_wait_finalize(self._ensure_ctx, timeout_s=float(timeout_s))
                return self._result
            self._result = self._task.result()
            return self._result

        op = await self.status()
        if not op.done():
            raise ReasonerNotReadyError(f"Operation not complete yet: operation_id={self.id} status={op.status}")

        if (op.status or "").upper() != OperationStatus.SUCCEEDED:
            raise ReasonerNotReadyError(
                op.error or f"Operation did not succeed: operation_id={self.id} status={op.status}"
            )

        self._result = await self._get(timeout_s=float(timeout_s))
        return self._result

    async def _get(self, *, timeout_s: float = 900) -> Reasoner:
        """Fetch and verify the reasoner after a successful operation."""
        name = self.metadata.get("reasoner_name")
        rtype = self.metadata.get("reasoner_type")
        if not isinstance(name, str) or not isinstance(rtype, str):
            raise ReasonerNotReadyError(
                "Operation completed but result metadata is missing "
                "(expected metadata: {'reasoner_name': str, 'reasoner_type': str})."
            )
        return await self.client._wait_until_ready(reasoner_name=name, reasoner_type=rtype, timeout_s=float(timeout_s))

    async def __aenter__(self) -> Reasoner:
        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"ReasonerOperation(id={self.id!r}, metadata={self.metadata!r})"


@dataclass
class ReasonerOperationSync:
    """Sync operation handle for reasoner lifecycle actions."""

    client: _ReasonersClientSyncProto
    id: str
    metadata: dict[str, Any]
    _result: Reasoner | None = None
    _ensure_requested: _EnsureRequested | None = None

    def status(self) -> Operation:
        """Return current operation status without waiting for READY."""
        return self.client.get_operation(self.id)

    def is_done(self) -> bool:
        """Return True if the operation is in a terminal state."""
        op = self.status()
        return OperationStatus.done(op.status)

    def wait(self, *, timeout_s: int = 900) -> Reasoner:
        """Wait for completion and return the resulting READY reasoner."""
        if self._result is not None:
            return self._result

        start = time.time()
        op = self.client.wait_for_operation(self.id, timeout_s=timeout_s)
        _raise_on_failed_operation(op, operation_id=self.id)

        remaining = timeout_s - int(max(0.0, time.time() - start))
        if remaining <= 0:
            raise ReasonerTimeoutError(f"Timed out after {timeout_s}s waiting for completion")
        self._result = self._get(timeout_s=remaining)
        return self._result

    def result(self, *, timeout_s: int = 900) -> Reasoner:
        """Return the operation result if complete (or raise if failed/incomplete)."""
        if self._result is not None:
            return self._result

        op = self.status()
        if not OperationStatus.done(op.status):
            raise ReasonerNotReadyError(f"Operation not complete yet: operation_id={self.id} status={op.status}")
        _raise_on_failed_operation(op, operation_id=self.id)

        self._result = self._get(timeout_s=timeout_s)
        return self._result

    def _get(self, *, timeout_s: int = 900) -> Reasoner:
        """Fetch and verify the reasoner after a successful operation."""
        name = self.metadata.get("reasoner_name")
        rtype = self.metadata.get("reasoner_type")
        if not isinstance(name, str) or not isinstance(rtype, str):
            raise ReasonerNotReadyError(
                "Operation completed but result metadata is missing "
                "(expected metadata: {'reasoner_name': str, 'reasoner_type': str})."
            )

        r = self.client.wait_until_ready(rtype, name, timeout_s=timeout_s)
        if self._ensure_requested is not None:
            self.client._raise_on_mismatch(
                reasoner_name=name,
                reasoner_type=rtype,
                requested=self._ensure_requested,
                existing=r,
            )
        return r

    def __enter__(self) -> Reasoner:
        return self.wait()

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

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


__all__ = [
    "ReasonerOperation",
    "ReasonerOperationSync",
    "_EnsureContext",
    "_EnsureRequested",
    "_op_from_reasoner",
    "_op_id",
    "_parse_op_id",
    "_wait_for_operation_async",
    "_wait_for_operation_sync",
]

