"""Client for interacting with a :class:`NatsJobService`."""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Any

from nats.aio.client import Client as NatsClient
from nats.aio.msg import Msg as CoreMsg
from nats.aio.subscription import Subscription
from pydantic import BaseModel, TypeAdapter

from nats_common.services.models import (
    JobEnvelope,
    JobStatus,
    JobStatusUpdate,
)
from nats_common.services.utils import _encode_payload, _extract_generic_args, _subject_with_id

logger = logging.getLogger(__name__)

_TERMINAL_STATUSES: frozenset[JobStatus] = frozenset(
    {
        JobStatus.FINISHED,
        JobStatus.CANCELLED,
        JobStatus.FAILED,
        JobStatus.TIMED_OUT,
    }
)


@dataclass
class JobCompletion[FinalT: BaseModel | None]:
    """Outcome returned by :meth:`NatsJobClient.wait_for_completion`.

    Attributes
    ----------
    status:
        The terminal :class:`JobStatusUpdate` that ended the job.
    final_result:
        The deserialized final result, or ``None`` if the service did not
        publish one (e.g. status-only jobs or failures).
    """

    status: JobStatusUpdate
    final_result: FinalT | None = None

    @property
    def succeeded(self) -> bool:
        """``True`` if the job finished successfully."""
        return self.status.status == JobStatus.FINISHED


class NatsJobClient[
    ReqT: BaseModel | None,
    FinalT: BaseModel | None,
    InterT: BaseModel | None,
]:
    """Client-side companion to :class:`~nats_common.services.jobs.NatsJobService`.

    Provides methods to submit jobs, request cancellation, subscribe to
    status / result subjects, and a high-level :meth:`wait_for_completion`
    convenience that bundles subscriptions and waits for a terminal status.

    The intermediate and final model types are automatically deduced from
    the generic type parameters — no need to pass model classes explicitly.

    Parameters
    ----------
    nats_client:
        An already-connected :class:`nats.aio.client.Client`.
    subject:
        Base subject matching the service (e.g. ``"myapp.ocr"``).

    Usage::

        client = NatsJobClient[MyRequest, MyChunk, MyResult](
            nats_client=nats_client,
            subject="myapp.ocr",
        )

        await client.submit(job_id="job-1", data=MyRequest(url="..."))
        completion = await client.wait_for_completion("job-1", timeout_seconds=60)

        if completion.succeeded:
            print(completion.final_result)
    """

    def __init__(
        self,
        *,
        nats_client: NatsClient,
        subject: str,
    ) -> None:
        self._nats_client = nats_client
        self.subject = subject.rstrip(".")

        # Same derived subjects as NatsJobService
        self.request_subject = f"{self.subject}.request"
        self.intermediate_result_subject = f"{self.subject}.intermediate_result"
        self.final_result_subject = f"{self.subject}.final_result"
        self.status_subject = f"{self.subject}.status"
        self.cancel_subject = f"{self.subject}.cancel"

        self._js = nats_client.jetstream()

        # Status adapter is always available (model is known statically).
        self._status_adapter = TypeAdapter(JobStatusUpdate)

        # Adapters for intermediate / final results are resolved lazily
        # because __orig_class__ is set *after* __init__ returns.
        self._adapters_resolved = False
        self._intermediate_adapter: TypeAdapter[Any] | None = None
        self._final_adapter: TypeAdapter[Any] | None = None

    # -- Lazy adapter resolution -------------------------------------------

    def _resolve_adapters(self) -> None:
        """Build :class:`TypeAdapter` instances from the generic type parameters.

        Called lazily on first access because ``__orig_class__`` is only
        available after ``__init__`` completes (Python sets it on the
        instance when the class is subscripted, e.g.
        ``NatsJobClient[Req, Inter, Final](...)``).
        """
        if self._adapters_resolved:
            return
        self._adapters_resolved = True

        args = _extract_generic_args(self, NatsJobClient)
        # args order: ReqT (index 0), InterT (index 1), FinalT (index 2)
        if len(args) >= 2:
            final_type = args[1]
            if final_type is not None and final_type is not type(None):
                self._final_adapter = TypeAdapter(final_type)
        if len(args) >= 3:
            inter_type = args[2]
            if inter_type is not None and inter_type is not type(None):
                self._intermediate_adapter = TypeAdapter(inter_type)

    # -- Job control -------------------------------------------------------

    async def submit(self, *, job_id: str, data: ReqT | None = None) -> None:
        """Submit a job request via JetStream publish."""
        envelope = JobEnvelope[Any](job_id=job_id, data=data)
        await self._js.publish(self.request_subject, _encode_payload(envelope))

    async def request_cancel(self, *, job_id: str) -> None:
        """Ask the service to cancel a running job (best-effort, core NATS)."""
        subj = _subject_with_id(self.cancel_subject, job_id)
        await self._nats_client.publish(subj, b"1")

    # -- Callback-based subscriptions --------------------------------------

    async def subscribe_status(
        self,
        job_id: str,
        callback: Callable[[JobStatusUpdate], Awaitable[None]],
    ) -> Subscription:
        """Subscribe to status updates for a single job.

        Returns the NATS :class:`Subscription` — call
        ``await sub.unsubscribe()`` when done.
        """
        subj = _subject_with_id(self.status_subject, job_id)

        async def _handler(msg: CoreMsg) -> None:
            update = self._status_adapter.validate_json(msg.data)

            # Check if callback is awaitable or not
            if asyncio.iscoroutinefunction(callback):
                await callback(update)
            else:
                callback(update)

        return await self._nats_client.subscribe(subj, cb=_handler)

    async def subscribe_status_all(
        self,
        callback: Callable[[JobStatusUpdate], Awaitable[None]],
    ) -> Subscription:
        """Subscribe to status updates for *all* jobs on this subject."""
        subj = f"{self.status_subject}.*"

        async def _handler(msg: CoreMsg) -> None:
            update = self._status_adapter.validate_json(msg.data)
            if asyncio.iscoroutinefunction(callback):
                await callback(update)
            else:
                callback(update)

        return await self._nats_client.subscribe(subj, cb=_handler)

    async def subscribe_intermediate_results(
        self,
        job_id: str,
        callback: Callable[[InterT], Awaitable[None]],
    ) -> Subscription:
        """Subscribe to intermediate results for a single job."""
        self._resolve_adapters()
        subj = _subject_with_id(self.intermediate_result_subject, job_id)

        async def _handler(msg: CoreMsg) -> None:
            if self._intermediate_adapter is not None:
                result = self._intermediate_adapter.validate_json(msg.data)
            else:
                result = msg.data  # type: ignore[assignment]
            if asyncio.iscoroutinefunction(callback):
                await callback(result)
            else:
                callback(result)

        return await self._nats_client.subscribe(subj, cb=_handler)

    async def subscribe_final_result(
        self,
        job_id: str,
        callback: Callable[[FinalT], Awaitable[None]],
    ) -> Subscription:
        """Subscribe to the final result for a single job."""
        self._resolve_adapters()
        subj = _subject_with_id(self.final_result_subject, job_id)

        async def _handler(msg: CoreMsg) -> None:
            if self._final_adapter is not None:
                result = self._final_adapter.validate_json(msg.data)
            else:
                result = msg.data  # type: ignore[assignment]
            await callback(result)

        return await self._nats_client.subscribe(subj, cb=_handler)

    # -- High-level convenience --------------------------------------------

    async def wait_for_completion(
        self,
        job_id: str,
        *,
        timeout_seconds: float | None = None,
        on_status: Callable[[JobStatusUpdate], Awaitable[None]] | None = None,
        on_intermediate: Callable[[InterT], Awaitable[None]] | None = None,
    ) -> JobCompletion[FinalT]:
        """Wait until the job reaches a terminal status.

        Subscribes to status, intermediate results (when *on_intermediate* is
        given) and the final result **before** returning, so no messages are
        missed when called right after :meth:`submit`.

        Parameters
        ----------
        job_id:
            The job to watch.
        timeout_seconds:
            Maximum seconds to wait.  ``None`` means wait indefinitely.
        on_status:
            Optional async callback invoked for every status update.
        on_intermediate:
            Optional async callback invoked for every intermediate result.

        Returns
        -------
        JobCompletion[FinalT]:
            Contains the terminal :class:`JobStatusUpdate` and, if published,
            the deserialized final result.

        Raises
        ------
        asyncio.TimeoutError:
            If *timeout_seconds* is exceeded before a terminal status arrives.
        """
        self._resolve_adapters()
        terminal_event = asyncio.Event()
        completion: JobCompletion[FinalT] = JobCompletion(
            status=JobStatusUpdate(job_id=job_id, status=JobStatus.PENDING),
        )
        subs: list[Subscription] = []

        try:
            # --- status subscription ---
            async def _on_status(update: JobStatusUpdate) -> None:
                completion.status = update
                if on_status is not None:
                    await on_status(update)
                if update.status in _TERMINAL_STATUSES:
                    terminal_event.set()

            subs.append(await self.subscribe_status(job_id, _on_status))

            # --- intermediate results subscription ---
            if on_intermediate is not None:
                subs.append(
                    await self.subscribe_intermediate_results(
                        job_id,
                        on_intermediate,
                    ),
                )

            # --- final result subscription ---
            print(self._final_adapter)
            if self._final_adapter is not None:

                async def _on_final(result: FinalT) -> None:
                    completion.final_result = result

                subs.append(
                    await self.subscribe_final_result(job_id, _on_final),
                )

            # --- wait for terminal status ---
            if timeout_seconds is not None:
                await asyncio.wait_for(
                    terminal_event.wait(),
                    timeout=timeout_seconds,
                )
            else:
                await terminal_event.wait()

            return completion

        finally:
            for sub in subs:
                try:
                    await sub.unsubscribe()
                except Exception:
                    logger.exception(
                        "Failed to unsubscribe during wait_for_completion",
                    )
