"""Reusable Job service primitives on top of NATS / JetStream.

Provides :class:`NatsJobService` for long-running, exactly-once jobs with
heartbeat, cancellation and progress reporting.

Accepts a plain :class:`nats.aio.client.Client` — no framework coupling.
"""

from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass
from enum import Enum
from typing import Any

from nats.aio.client import Client as NatsClient
from nats.aio.msg import Msg as CoreMsg
from nats.js.api import (
    AckPolicy,
    ConsumerConfig,
    DeliverPolicy,
    RetentionPolicy,
    StorageType,
    StreamConfig,
)
from nats.js.errors import BadRequestError, FetchTimeoutError
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, _sanitize_name, _subject_with_id

logger = logging.getLogger(__name__)


# ===============================================================================
# region Configuration
# ===============================================================================


class PublishMode(str, Enum):
    """Whether to publish via core NATS or JetStream."""

    CORE = "core"
    JETSTREAM = "jetstream"


@dataclass(frozen=True)
class JobPublishConfig:
    """Controls how job status / result / cancel messages are published."""

    status: PublishMode = PublishMode.CORE
    intermediate_result: PublishMode = PublishMode.CORE
    final_result: PublishMode = PublishMode.CORE
    cancel: PublishMode = PublishMode.CORE

    # Only relevant when at least one mode above is JETSTREAM.
    publish_stream_name: str | None = None
    publish_stream_max_age_seconds: float = 60 * 60  # 1 h
    publish_stream_storage: StorageType = StorageType.FILE
    publish_stream_max_msgs: int | None = None


@dataclass(frozen=True)
class JobQueueConfig:
    """Controls the JetStream work-queue stream and consumer for job requests."""

    stream_name: str | None = None
    consumer_name: str | None = None
    stream_storage: StorageType = StorageType.FILE
    stream_max_age_seconds: float = 7 * 24 * 3600  # 7 d
    consumer_ack_wait_seconds: float = 60.0
    consumer_max_ack_pending: int | None = None  # None -> derived from num_max_parallel_jobs


# ===============================================================================
# region Helpers
# ===============================================================================


def _extract_request_type(
    job_cls: type[NatsJob[Any, Any, Any]],
) -> type[BaseModel] | None:
    """Extract the ``ReqT`` type parameter from a :class:`NatsJob` subclass.

    Uses :func:`_extract_generic_args` to find the concrete type argument
    supplied to ``NatsJob[ReqT, FinalT, InterT]``.  Returns ``None`` when the
    request type is ``None`` / ``NoneType`` (i.e. jobs that carry no payload).
    """
    args = _extract_generic_args(job_cls, NatsJob)
    if args:
        req_type = args[0]
        if req_type is None or req_type is type(None):
            return None
        return req_type  # type: ignore[return-value]
    return None


# ===============================================================================
# region NatsJobService
# ===============================================================================


class NatsJobService[
    ReqT: BaseModel | None,
    InterT: BaseModel | None,
    FinalT: BaseModel | None,
]:
    """Server-side orchestrator for long-running jobs.

    * Requests are consumed exactly-once via a JetStream **work-queue** stream.
    * Status, intermediate and final results are published on per-``job_id``
      subjects.
    * Cancellation is delivered via a lightweight core-NATS subscription.

    The request model for envelope validation is automatically deduced from
    the type parameters of the supplied ``job_cls``.

    Parameters
    ----------
    nats_client:
        An already-connected :class:`nats.aio.client.Client`.
        If using FastStream, use the `broker._connection` attribute.
    subject:
        Base subject (e.g. ``"myapp.ocr"``).  Derived subjects like
        ``myapp.ocr.request``, ``myapp.ocr.status.<job_id>`` etc. are created
        automatically.
    job_cls:
        Concrete :class:`NatsJob` subclass whose :meth:`~NatsJob.handle`
        method contains the actual work.  The ``ReqT`` type parameter is
        extracted at runtime to build the request envelope validator.
    num_max_parallel_jobs:
        Maximum number of jobs that can run in parallel.
    heartbeat_interval_seconds:
        Interval at which to publish heartbeat status updates.
    publish_config:
        Configuration for publishing job status, intermediate and final results.
    queue_config:
        Configuration for the JetStream work-queue stream and consumer.
    fetch_timeout_seconds:
        Timeout for fetching job requests from the JetStream work-queue stream.

    Usage::

        class MyJob(NatsJob[MyRequest, MyResult, MyPartial]):
            async def handle(self) -> None:
                ...

        svc = NatsJobService(
            nats_client=nats_client,
            subject="myapp.ocr",
            job_cls=MyJob,
        )
        await svc.start()
    """

    def __init__(
        self,
        *,
        nats_client: NatsClient,
        subject: str,
        job_cls: type[NatsJob[ReqT, FinalT, InterT]],
        num_max_parallel_jobs: int = 1,
        heartbeat_interval_seconds: float = 5.0,
        publish_config: JobPublishConfig = JobPublishConfig(),
        queue_config: JobQueueConfig = JobQueueConfig(),
        fetch_timeout_seconds: float = 1.0,
    ) -> None:
        self._nats_client = nats_client
        self.subject = subject.rstrip(".")

        self.job_cls = job_cls

        self.num_max_parallel_jobs = num_max_parallel_jobs
        self.heartbeat_interval_seconds = heartbeat_interval_seconds
        self.publish_cfg = publish_config
        self.queue_cfg = queue_config
        self.fetch_timeout_seconds = fetch_timeout_seconds

        # Derived subjects
        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"

        # JetStream
        self._js = nats_client.jetstream()
        self._jsm: Any | None = None

        # Subscriptions & tasks
        self._pull_sub: Any | None = None
        self._cancel_sub: Any | None = None
        self._stop = asyncio.Event()
        self._pull_task: asyncio.Task[None] | None = None

        # Concurrency control
        self._sema = asyncio.Semaphore(num_max_parallel_jobs)
        self._running: dict[str, NatsJob[ReqT, FinalT, InterT]] = {}
        self._running_tasks: dict[str, asyncio.Task[None]] = {}

        # Pydantic adapter for envelope validation — request type is
        # automatically extracted from the job_cls type parameters.
        request_model = _extract_request_type(job_cls)
        if request_model is None:
            self._request_adapter: TypeAdapter[Any] = TypeAdapter(JobEnvelope[Any])
        else:
            self._request_adapter = TypeAdapter(
                JobEnvelope[request_model],  # type: ignore[valid-type]
            )

    # -- Lifecycle ---------------------------------------------------------

    async def start(self) -> None:
        """Provision JetStream resources and start consuming job requests."""
        self._stop.clear()
        self._jsm = self._nats_client.jsm()

        await self._ensure_job_queue()
        await self._ensure_publish_streams()

        self._pull_sub = await self._js.pull_subscribe_bind(
            stream=self._stream_name(),
            consumer=self._consumer_name(),
        )
        self._cancel_sub = await self._nats_client.subscribe(
            f"{self.cancel_subject}.*",
            cb=self._on_cancel,
        )
        self._pull_task = asyncio.create_task(
            self._pull_loop(),
            name=f"jobsvc:{self.subject}:pull",
        )

    async def stop(self) -> None:
        """Gracefully stop the service, cancelling all running jobs."""
        self._stop.set()

        if self._pull_task and not self._pull_task.done():
            self._pull_task.cancel()

        for task in list(self._running_tasks.values()):
            if not task.done():
                task.cancel()

        for sub, label in (
            (self._cancel_sub, "cancel"),
            (self._pull_sub, "pull"),
        ):
            if sub is not None:
                try:
                    await sub.unsubscribe()
                except Exception:
                    logger.exception("Failed to unsubscribe %s subscription", label)

    # -- Client helpers ----------------------------------------------------

    async def submit(self, *, job_id: str, data: ReqT | None = None) -> None:
        """Enqueue a new 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:
        """Request cancellation of a running job (best-effort)."""
        subj = _subject_with_id(self.cancel_subject, job_id)
        await self._publish(subj, b"1", mode=self.publish_cfg.cancel)

    # -- Internal: naming --------------------------------------------------

    def _stream_name(self) -> str:
        return self.queue_cfg.stream_name or _sanitize_name(f"{self.subject}_JOBQ")

    def _consumer_name(self) -> str:
        return self.queue_cfg.consumer_name or _sanitize_name(f"{self.subject}_JOBQ_CONSUMER")

    def _publish_stream_name(self) -> str:
        return self.publish_cfg.publish_stream_name or _sanitize_name(
            f"{self.subject}_JOBPUB",
        )

    # -- Internal: provisioning --------------------------------------------

    async def _ensure_job_queue(self) -> None:
        assert self._jsm is not None
        stream = self._stream_name()

        await self._jsm.add_stream(
            config=StreamConfig(
                name=stream,
                subjects=[self.request_subject],
                retention=RetentionPolicy.WORK_QUEUE,
                storage=self.queue_cfg.stream_storage,
                max_age=self.queue_cfg.stream_max_age_seconds,
            ),
        )

        max_ack_pending = self.queue_cfg.consumer_max_ack_pending if self.queue_cfg.consumer_max_ack_pending is not None else max(1, self.num_max_parallel_jobs * 2)

        try:
            await self._jsm.add_consumer(
                stream=stream,
                config=ConsumerConfig(
                    durable_name=self._consumer_name(),
                    ack_policy=AckPolicy.EXPLICIT,
                    ack_wait=self.queue_cfg.consumer_ack_wait_seconds,
                    max_ack_pending=max_ack_pending,
                    filter_subject=self.request_subject,
                    deliver_policy=DeliverPolicy.ALL,
                ),
            )
        except BadRequestError as exc:
            raise RuntimeError(
                f"JetStream consumer {self._consumer_name()!r} already exists with mismatched config: {exc}",
            ) from exc

    async def _ensure_publish_streams(self) -> None:
        assert self._jsm is not None

        modes = (
            self.publish_cfg.status,
            self.publish_cfg.intermediate_result,
            self.publish_cfg.final_result,
            self.publish_cfg.cancel,
        )
        if all(m != PublishMode.JETSTREAM for m in modes):
            return

        subjects: list[str] = []
        for mode, subj in (
            (self.publish_cfg.status, self.status_subject),
            (self.publish_cfg.intermediate_result, self.intermediate_result_subject),
            (self.publish_cfg.final_result, self.final_result_subject),
            (self.publish_cfg.cancel, self.cancel_subject),
        ):
            if mode == PublishMode.JETSTREAM:
                subjects.append(f"{subj}.*")

        await self._jsm.add_stream(
            config=StreamConfig(
                name=self._publish_stream_name(),
                subjects=subjects,
                retention=RetentionPolicy.LIMITS,
                storage=self.publish_cfg.publish_stream_storage,
                max_age=self.publish_cfg.publish_stream_max_age_seconds,
                max_msgs=self.publish_cfg.publish_stream_max_msgs,
            ),
        )

    # -- Internal: publishing ----------------------------------------------

    async def _publish(
        self,
        subject: str,
        payload: bytes,
        *,
        mode: PublishMode,
    ) -> None:
        if mode == PublishMode.JETSTREAM:
            await self._js.publish(subject, payload)
        else:
            await self._nats_client.publish(subject, payload)

    # -- Internal: cancel handler ------------------------------------------

    async def _on_cancel(self, msg: CoreMsg) -> None:
        parts = (msg.subject or "").split(".")
        if not parts:
            return
        job_id = parts[-1]
        job = self._running.get(job_id)
        if job is not None:
            await job.cancel()

    # -- Internal: pull loop & processing ----------------------------------

    async def _pull_loop(self) -> None:
        assert self._pull_sub is not None

        while not self._stop.is_set():
            # Block until there is capacity.
            try:
                await asyncio.wait_for(self._sema.acquire(), timeout=0.25)
            except asyncio.TimeoutError:
                continue

            try:
                msgs = await self._pull_sub.fetch(
                    1,
                    timeout=self.fetch_timeout_seconds,
                )
            except (FetchTimeoutError, TimeoutError):
                self._sema.release()
                continue
            except Exception:
                self._sema.release()
                logger.exception("Pull fetch failed")
                await asyncio.sleep(0.2)
                continue

            if not msgs:
                self._sema.release()
                continue

            # Hand off to a task; semaphore is released inside _process_one.
            asyncio.create_task(
                self._process_one(msgs[0]),
                name=f"job:{self.subject}",
            )

    async def _process_one(self, msg: Any) -> None:
        job_id: str | None = None
        try:
            envelope: JobEnvelope[Any] = self._request_adapter.validate_json(
                msg.data,
            )
            job_id = envelope.job_id

            job = self.job_cls(service=self, request=envelope, js_msg=msg)
            self._running[job_id] = job

            await job.update_status(JobStatus.STARTED, message="accepted")

            task = asyncio.create_task(
                job._run(),
                name=f"job:{self.subject}:{job_id}",
            )
            self._running_tasks[job_id] = task
            await task

            await msg.ack()

        except asyncio.CancelledError:
            try:
                await msg.ack()
            except Exception:
                pass
            raise

        except Exception as exc:
            logger.exception("Job processing failed (job_id=%s)", job_id)
            if job_id is not None:
                running_job = self._running.get(job_id)
                if running_job is not None:
                    try:
                        await running_job.update_status(
                            JobStatus.FAILED,
                            error=str(exc),
                        )
                    except Exception:
                        logger.exception("Failed to publish FAILED status")
            try:
                await msg.ack()
            except Exception:
                pass

        finally:
            if job_id is not None:
                self._running.pop(job_id, None)
                self._running_tasks.pop(job_id, None)
            self._sema.release()


# ===============================================================================
# region NatsJob
# ===============================================================================


class NatsJob[
    ReqT: BaseModel | None,
    FinalT: BaseModel | None,
    InterT: BaseModel | None,
]:
    """Base class for job handlers.

    Subclass and override :meth:`handle` to implement your job logic.
    Inside ``handle()`` you can:
    * get the request via ``self.request``
    * ``await self.publish_intermediate_result(...)`` - stream partial results.
    * ``await self.publish_final_result(...)`` - send the final result.
    * ``await self.update_status(...)`` - push an explicit status update.
    * ``self.set_local_status(...)`` - update the status that the heartbeat
      emits *without* publishing immediately.
    * Check ``self.cancel_requested`` to react to cancellation.
    """

    def __init__(
        self,
        *,
        service: NatsJobService[ReqT, InterT, FinalT],
        request: JobEnvelope[ReqT],
        js_msg: Any,
    ) -> None:
        self.service = service
        self.request = request
        self.job_id: str = request.job_id

        self._msg = js_msg
        self._cancel_evt = asyncio.Event()
        self._done_evt = asyncio.Event()

        self._last_status = JobStatusUpdate(
            job_id=self.job_id,
            status=JobStatus.PENDING,
            message="pending",
        )
        self._hb_task: asyncio.Task[None] | None = None

    # -- Status & heartbeat ------------------------------------------------

    @property
    def cancel_requested(self) -> bool:
        """``True`` once cancellation has been requested for this job."""
        return self._cancel_evt.is_set()

    def set_local_status(
        self,
        status: JobStatus,
        *,
        message: str | None = None,
        error: str | None = None,
        progress: float | None = None,
        estimated_time_remaining_seconds: float | None = None,
    ) -> None:

        if status == JobStatus.FINISHED:
            progress = 1.0
            estimated_time_remaining_seconds = 0.0

        """Update internal status (the heartbeat loop will broadcast it)."""
        self._last_status = JobStatusUpdate(
            job_id=self.job_id,
            status=status,
            message=message,
            error=error,
            progress=progress,
            estimated_time_remaining_seconds=estimated_time_remaining_seconds,
        )

    async def update_status(
        self,
        status: JobStatus,
        message: str | None = None,
        error: str | None = None,
        progress: float | None = None,
        estimated_time_remaining_seconds: float | None = None,
    ) -> None:
        """Update and *immediately* publish the job status."""
        self.set_local_status(
            status,
            message=message,
            error=error,
            progress=progress,
            estimated_time_remaining_seconds=estimated_time_remaining_seconds,
        )
        await self._publish_status()

    async def _publish_status(self) -> None:
        subj = _subject_with_id(self.service.status_subject, self.job_id)
        await self.service._publish(
            subj,
            _encode_payload(self._last_status),
            mode=self.service.publish_cfg.status,
        )

    async def _heartbeat_loop(self) -> None:
        """Periodically extend ack deadline and re-publish current status."""
        try:
            while not self._done_evt.is_set():
                # Extend the JetStream ack deadline so the message is not
                # redelivered while the job is still running.
                try:
                    await self._msg.in_progress()
                except Exception:
                    pass  # not a JetStream msg or already acked

                try:
                    await self._publish_status()
                except Exception:
                    logger.exception("Failed to publish heartbeat status")

                try:
                    await asyncio.wait_for(
                        self._done_evt.wait(),
                        timeout=self.service.heartbeat_interval_seconds,
                    )
                except asyncio.TimeoutError:
                    continue
        except asyncio.CancelledError:
            pass

    # -- Results -----------------------------------------------------------

    async def publish_intermediate_result(self, result: InterT) -> None:
        """Publish a partial / intermediate result for this job."""
        subj = _subject_with_id(
            self.service.intermediate_result_subject,
            self.job_id,
        )
        await self.service._publish(
            subj,
            _encode_payload(result),
            mode=self.service.publish_cfg.intermediate_result,
        )

    async def publish_final_result(self, result: FinalT) -> None:
        """Publish the final result for this job."""
        subj = _subject_with_id(self.service.final_result_subject, self.job_id)
        await self.service._publish(
            subj,
            _encode_payload(result),
            mode=self.service.publish_cfg.final_result,
        )

    # -- Cancellation ------------------------------------------------------

    async def cancel(self) -> None:
        """Signal this job to cancel (cooperative)."""
        self._cancel_evt.set()

    # -- Execution ---------------------------------------------------------

    async def handle(self) -> None:
        """Override in a subclass to implement the actual job logic."""
        raise NotImplementedError

    async def _run(self) -> None:
        self._hb_task = asyncio.create_task(
            self._heartbeat_loop(),
            name=f"job-hb:{self.service.subject}:{self.job_id}",
        )
        try:
            await self.update_status(JobStatus.RUNNING, message="running")
            await self.handle()

            if self.cancel_requested:
                await self.update_status(JobStatus.CANCELLED, message="cancelled")
                return

            await self.update_status(JobStatus.FINISHED, message="finished")

        except asyncio.CancelledError:
            await self.update_status(JobStatus.CANCELLED, message="cancelled")
            raise

        except Exception as exc:
            await self.update_status(JobStatus.FAILED, error=str(exc))
            raise

        finally:
            self._done_evt.set()
            if self._hb_task and not self._hb_task.done():
                self._hb_task.cancel()
