"""Async Jobs service client.

Exposed via `client.jobs` on the async-first root `Client`.

This service is a unified abstraction over two backend concepts:

- **Jobs** for `reasoner_type` in {`Prescriptive`, `Predictive`} (SQL or Direct Access).
- **Transactions** for `reasoner_type == Logic` (exposed as "jobs" to the user).

Backend capability notes (by design):

- **Logic list/cancel**: implemented via SQL against `{app}.api.transactions` and
  `API.CANCEL_(OWN_)TRANSACTION` (even when the configured backend is Direct Access).
- **Logic create**: prefers Direct Access when available; falls back to SQL (`{app}.api.exec_async_v2`).
- **Logic get/events**: prefer Direct Access when available, fall back to SQL.
"""

from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass
from datetime import datetime
from types import TracebackType
from typing import Any, TYPE_CHECKING

from relationalai.util.docutils import include_in_docs

from .constants import JOB_TERMINAL_STATES, JOB_TYPES, MAX_POLL_INTERVAL_S
from .models import Job, JobEvents, record_to_job_with_include
from .gateways import JobsGateway
from .errors import JobInvalidRequestError, JobTimeoutError
from .models import Operation
from .operations import JobOperation, operation_from_job
from .util import (
    JobIncludeDict,
    normalize_abort_detached_payload,
    normalize_include,
    normalize_optional_nonempty_str,
    normalize_public_reasoner_type,
    normalize_public_reasoner_type_label,
    raise_on_client_disconnected,
    require_nonempty_str,
    require_payload_dict,
    validate_include_for_type,
)

if TYPE_CHECKING:
    from relationalai.config.config import Config

_LOG = logging.getLogger(__name__)

__include_in_docs__ = True


async def _fetch_expansions_async(
    gw: JobsGateway,
    *,
    reasoner_type: str,
    job_id: str,
    include: set[str],
) -> tuple[JobEvents | None, dict[str, Any] | None, dict[str, Any] | None]:
    """Best-effort fetch of expanded fields for a single job/txn."""

    async def _events() -> JobEvents:
        return await gw.events(reasoner_type, job_id, continuation_token="", stream_name="progress")

    async def _artifacts() -> dict[str, Any] | None:
        return await gw.artifacts(reasoner_type, job_id)

    async def _problems() -> dict[str, Any] | None:
        return await gw.problems(reasoner_type, job_id)

    events_task = _events() if "events" in include else None
    artifacts_task = _artifacts() if "artifacts" in include else None
    problems_task = _problems() if "problems" in include else None

    coros = [c for c in (events_task, artifacts_task, problems_task) if c is not None]
    if not coros:
        return None, None, None

    # Gather concurrently; never fail the whole view because an expansion failed.
    results = await asyncio.gather(*coros, return_exceptions=True)
    idx = 0

    events: JobEvents | None = None
    if events_task is not None:
        r = results[idx]
        idx += 1
        events = JobEvents(events=[], continuation_token=None) if isinstance(r, BaseException) else r  # type: ignore[assignment]

    artifacts: dict[str, Any] | None = None
    if artifacts_task is not None:
        r = results[idx]
        idx += 1
        artifacts = {} if isinstance(r, BaseException) else (r or {})  # type: ignore[assignment]

    problems: dict[str, Any] | None = None
    if problems_task is not None:
        r = results[idx]
        problems = {} if isinstance(r, BaseException) else (r or {})  # type: ignore[assignment]

    return events, artifacts, problems


@include_in_docs
@dataclass
class JobsClient:
    """Async-first jobs client.

    API behavior depends on the public `reasoner_type` value:
    - `Logic` routes to transactions.
    - `Prescriptive` / `Predictive` route to jobs.
    """

    _gateway: JobsGateway
    _config: "Config | None" = None
    _closed: bool = False

    @include_in_docs
    async def aclose(self) -> None:
        """Close underlying resources (best-effort).

        Note: SQL gateways typically have nothing to close; Direct Access gateways may
        own an async HTTP client/session.

        Returns
        -------
        None
            Always returns `None`.
        """
        await self._gateway.aclose()
        self._closed = True

    def __del__(self) -> None:
        # No transport ownership here; keep as best-effort warning for future async gateways.
        try:
            if getattr(self, "_closed", True):
                return
            _LOG.debug("JobsClient was not closed (no-op for SQL gateways).")
        except Exception:
            return

    async def __aenter__(self) -> "JobsClient":
        return self

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

    @include_in_docs
    async def get(self, reasoner_type: str, job_id: str, *, include: JobIncludeDict | None = None) -> Job | None:
        """Get a normalized job/transaction by `(reasoner_type, job_id)`.

        By default returns a stable common schema. Use `include=` to opt into
        backend-specific expansions (e.g. `details`, `events`, `artifacts`, `problems`, `raw`).

        Returns
        -------
        Job | None
            The job/transaction, or `None` if not found.
        """
        t = normalize_public_reasoner_type(reasoner_type, field="reasoner_type")
        jid = require_nonempty_str("job_id", job_id)
        return await self._get_canonical(t, jid, include=include)

    async def _get_canonical(self, reasoner_type: str, job_id: str, *, include: JobIncludeDict | None) -> Job | None:
        """Internal get() that expects canonical Jobs discriminators."""
        t = require_nonempty_str("reasoner_type", reasoner_type)
        jid = require_nonempty_str("job_id", job_id)
        inc = normalize_include(include)
        validate_include_for_type(t, inc)
        job = await self._gateway.get(t, jid)
        if job is None:
            return None

        events, artifacts, problems = await _fetch_expansions_async(
            self._gateway, reasoner_type=t, job_id=jid, include=inc
        )
        return record_to_job_with_include(job, include=inc, events=events, artifacts=artifacts, problems=problems)

    # --------------------------------------------------
    # Job API
    # --------------------------------------------------
    @include_in_docs
    async def create(
        self,
        reasoner_type: str,
        reasoner_name: str,
        payload: dict[str, Any],
        *,
        timeout_mins: int | None = None,
    ) -> JobOperation:
        """Start a job/transaction and return a non-blocking operation handle.

        The returned handle's `id` is the `job_id`/`txn_id` to poll with `wait(...)` /
        `wait_for_operation(...)`.

        Note: `reasoner_name` identifies which reasoner/worker to run on.

        Returns
        -------
        JobOperation
            An operation handle; call `await op.wait(...)` to block for completion.
        """
        public_rt = require_nonempty_str("reasoner_type", reasoner_type)
        public_label = normalize_public_reasoner_type_label(public_rt, field="reasoner_type")
        t = normalize_public_reasoner_type(public_rt, field="reasoner_type")
        n = require_nonempty_str("reasoner_name", reasoner_name)
        payload_d = require_payload_dict(payload)
        safe_payload = normalize_abort_detached_payload(reasoner_type=t, payload=payload_d, enable=False)
        id = await self._gateway.create(t, n, safe_payload, timeout_mins=timeout_mins)
        # Store public reasoner type name in operation metadata; internal waits/statuses
        # should never require callers to know Jobs gateway discriminators (SOLVER/ML).
        return JobOperation(client=self, id=id, metadata={"reasoner_type": public_label, "name": n})

    @include_in_docs
    async def create_ready(
        self,
        reasoner_type: str,
        reasoner_name: str,
        payload: dict[str, Any],
        *,
        timeout_mins: int | None = None,
        timeout_s: float = 900,
        poll_interval_s: float = 0.5,
        include: JobIncludeDict | None = None,
    ) -> Job:
        """Create a job and block until it reaches a terminal state.

        Returns
        -------
        Job
            The final job/transaction (possibly expanded when `include=` is provided).
        """
        public_rt = require_nonempty_str("reasoner_type", reasoner_type)
        public_label = normalize_public_reasoner_type_label(public_rt, field="reasoner_type")
        t = normalize_public_reasoner_type(public_rt, field="reasoner_type")
        n = require_nonempty_str("reasoner_name", reasoner_name)
        payload_d = require_payload_dict(payload)
        safe_payload = normalize_abort_detached_payload(reasoner_type=t, payload=payload_d, enable=True)
        id = await self._gateway.create(t, n, safe_payload, timeout_mins=timeout_mins)
        op = JobOperation(client=self, id=id, metadata={"reasoner_type": public_label, "name": n})
        return await op.wait(timeout_s=timeout_s, poll_interval_s=poll_interval_s, include=include)

    @include_in_docs
    async def list(
        self,
        reasoner_type: str | None = None,
        *,
        job_id: str | None = None,
        state: str | None = None,
        name: str | None = None,
        created_by: str | None = None,
        only_active: bool = False,
        all_users: bool = False,
        limit: int | None = None,
        include: JobIncludeDict | None = None,
    ) -> list[Job]:
        """List jobs/transactions, optionally filtered and optionally expanded.

        - If `reasoner_type` is `None`, this lists across all supported job types
          and merges the results, sorting by `created_on` (desc).
        - `include=` enables best-effort expansions per returned item.

        Returns
        -------
        list[Job]
            Matching jobs/transactions (possibly expanded when `include=` is provided).
        """
        if reasoner_type is None:
            ts = sorted(JOB_TYPES)
            all_items: list[Job] = []
            for t in ts:
                all_items.extend(
                    await self._list_canonical(
                        t,
                        job_id=job_id,
                        state=state,
                        name=name,
                        created_by=created_by,
                        only_active=only_active,
                        all_users=all_users,
                        limit=None,
                        include=include,
                    )
                )

            def _key(j: Job) -> datetime:
                v = getattr(j, "created_on", None)
                return v if isinstance(v, datetime) else datetime.min

            all_items.sort(key=_key, reverse=True)
            if limit is not None and int(limit) > 0:
                return all_items[: int(limit)]
            return all_items

        t = normalize_public_reasoner_type(reasoner_type, field="reasoner_type")
        return await self._list_canonical(
            t,
            job_id=job_id,
            state=state,
            name=name,
            created_by=created_by,
            only_active=only_active,
            all_users=all_users,
            limit=limit,
            include=include,
        )

    async def _list_canonical(
        self,
        reasoner_type: str,
        *,
        job_id: str | None = None,
        state: str | None = None,
        name: str | None = None,
        created_by: str | None = None,
        only_active: bool = False,
        all_users: bool = False,
        limit: int | None = None,
        include: JobIncludeDict | None = None,
    ) -> list[Job]:
        """Internal list() that expects canonical Jobs discriminators."""
        t = require_nonempty_str("reasoner_type", reasoner_type)
        jid = normalize_optional_nonempty_str("job_id", job_id)
        st = normalize_optional_nonempty_str("state", state)
        n = normalize_optional_nonempty_str("name", name)
        cb = normalize_optional_nonempty_str("created_by", created_by)
        if not isinstance(only_active, bool):
            raise JobInvalidRequestError("Job `only_active` must be a bool.")
        if not isinstance(all_users, bool):
            raise JobInvalidRequestError("Job `all_users` must be a bool.")
        if limit is not None and int(limit) <= 0:
            return []
        inc = normalize_include(include)
        validate_include_for_type(t, inc)
        jobs = await self._gateway.list(
            t,
            job_id=jid,
            state=st,
            name=n,
            created_by=cb,
            only_active=only_active,
            all_users=all_users,
            limit=limit,
        )

        out: list[Job] = []
        for j in jobs:
            events, artifacts, problems = await _fetch_expansions_async(
                self._gateway, reasoner_type=t, job_id=j.id, include=inc
            )
            out.append(record_to_job_with_include(j, include=inc, events=events, artifacts=artifacts, problems=problems))

        return out

    @include_in_docs
    async def get_events(
        self,
        reasoner_type: str,
        job_id: str,
        continuation_token: str = "",
        *,
        stream_name: str = "progress",
    ) -> JobEvents:
        """Get events for a job or transaction.

        Works for all reasoner types and backends:
        - LOGIC (DA): ``GET /v1alpha1/transactions/{txn_id}/events/{stream_name}``
        - LOGIC (SQL): ``{app}.api.get_own_transaction_events(txn_id, continuation_token)``
        - Non-LOGIC (DA): ``GET /v1alpha1/jobs/{job_type}/{job_id}/events/{stream_name}``
        - Non-LOGIC (SQL): ``{app}.{schema}.GET_JOB_EVENTS(reasoner_type, job_id, token)``

        Returns
        -------
        JobEvents
            Event page including `events` and an optional `continuation_token`.
        """
        t = normalize_public_reasoner_type(reasoner_type, field="reasoner_type")
        jid = require_nonempty_str("job_id", job_id)
        if not isinstance(continuation_token, str):
            raise JobInvalidRequestError("Job `continuation_token` must be a string.")
        if not isinstance(stream_name, str):
            raise JobInvalidRequestError("Job `stream_name` must be a string.")
        return await self._gateway.events(t, jid, continuation_token=continuation_token, stream_name=stream_name)

    @include_in_docs
    async def cancel(self, reasoner_type: str, job_id: str, *, all_users: bool = False) -> None:
        """Cancel a running job/transaction.

        Cancellation is supported for both non-LOGIC jobs and LOGIC transactions.

        `cancel(all_users=True)` is only supported for LOGIC transactions; non-LOGIC
        jobs reject it.

        Returns
        -------
        None
            Always returns `None`.
        """
        t = normalize_public_reasoner_type(reasoner_type, field="reasoner_type")
        jid = require_nonempty_str("job_id", job_id)
        if not isinstance(all_users, bool):
            raise JobInvalidRequestError("Job `all_users` must be a bool.")
        return await self._gateway.cancel(t, jid, all_users=all_users)

    @include_in_docs
    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:
        """Poll `get(...)` until the execution reaches a terminal state.

        Returns the final `Job` (or `JobExpanded` if `include=` is provided).
        Raises `JobTimeoutError` when `timeout_s` is exceeded.

        Returns
        -------
        Job
            The final job/transaction (possibly expanded when `include=` is provided).
        """
        t = normalize_public_reasoner_type(reasoner_type, field="reasoner_type")
        jid = require_nonempty_str("job_id", job_id)
        if timeout_s is None or float(timeout_s) <= 0:
            raise JobInvalidRequestError("Job `timeout_s` must be > 0.")
        if poll_interval_s is None or float(poll_interval_s) <= 0:
            raise JobInvalidRequestError("Job `poll_interval_s` must be > 0.")
        if float(poll_interval_s) >= MAX_POLL_INTERVAL_S:
            raise JobInvalidRequestError(f"Job `poll_interval_s` must be < {int(MAX_POLL_INTERVAL_S)} seconds.")

        interval_s = min(float(poll_interval_s), MAX_POLL_INTERVAL_S - 1e-9)
        deadline = time.monotonic() + float(timeout_s)
        while True:
            v = await self._get_canonical(t, jid, include=include)
            if v is not None:
                st = str(v.state or "").upper()
                if st in JOB_TERMINAL_STATES:
                    raise_on_client_disconnected(reasoner_type=t, job=v)
                    return v

            now = time.monotonic()
            if now >= deadline:
                raise JobTimeoutError(f"Timed out after {timeout_s}s waiting for completion: reasoner_type={t} job_id={jid}")
            await asyncio.sleep(min(interval_s, max(0.0, deadline - now)))

    @include_in_docs
    async def get_operation(self, reasoner_type: str, job_id: str) -> Operation:
        """Derive a normalized operation status from the current job state.

        Returns
        -------
        Operation
            A normalized operation status.
        """
        t = normalize_public_reasoner_type(reasoner_type, field="reasoner_type")
        jid = require_nonempty_str("job_id", job_id)
        job = await self._get_canonical(t, jid, include=None)
        return operation_from_job(jid, job)



__all__ = ["JobsClient"]

