"""Async worker service for background job execution.

The ``AsyncWorker`` service polls for queued jobs and executes them
concurrently via a ``ThreadPoolExecutor``.  It supports configurable
concurrency, graceful shutdown via signal handling, stuck job detection,
and job cleanup.

## Features

- Polling loop with configurable interval (``async.poll_interval``)
- True concurrent execution via ThreadPoolExecutor (``async.max_workers``)
- Graceful shutdown via SIGINT/SIGTERM signal handling
- Stuck job detection: mark jobs failed after TTL if no heartbeat
- Worker health report (last_heartbeat, jobs processed / failed / cancelled)
- Cancellation token support with propagation to tool execution
- Job cleanup routine to prune completed jobs older than retention
- Race-safe cancellation: ``cancel_job`` only signals the token;
  ``pickup_and_execute`` owns all state transitions

## Configuration Keys

| Key                    | Type  | Default | Description                        |
|------------------------|-------|---------|------------------------------------|
| ``async.enabled``      | bool  | False   | Enable async job execution         |
| ``async.max_workers``  | int   | 4       | Max concurrent worker threads      |
| ``async.poll_interval``| float | 1.0     | Poll interval in seconds           |
| ``async.job_timeout``  | int   | 3600    | Job timeout in seconds             |
| ``async.job_ttl``      | int   | 86400   | Completed job retention (sec)      |

Based on issue #312 (feat(async): add async command execution and workers).
"""

from __future__ import annotations

import logging
import signal
import threading
from concurrent.futures import Future, ThreadPoolExecutor
from datetime import UTC, datetime, timedelta
from typing import Any
from uuid import uuid4

from cleveragents.domain.models.core.async_job import (
    AsyncJob,
    AsyncJobStatus,
    InvalidJobTransitionError,
)
from cleveragents.shared.redaction import redact_value

logger = logging.getLogger(__name__)


class CancellationToken:
    """Thread-safe cancellation token for cooperative cancellation.

    Workers check this token periodically and stop execution when
    cancellation is requested.
    """

    def __init__(self) -> None:
        self._cancelled = threading.Event()

    @property
    def is_cancelled(self) -> bool:
        """Check if cancellation has been requested."""
        return self._cancelled.is_set()

    def cancel(self) -> None:
        """Request cancellation."""
        self._cancelled.set()

    def reset(self) -> None:
        """Reset the cancellation token."""
        self._cancelled.clear()

    def wait(self, timeout: float | None = None) -> bool:
        """Wait for cancellation or timeout.

        Args:
            timeout: Maximum seconds to wait (None = wait forever).

        Returns:
            True if cancelled, False if timed out.
        """
        return self._cancelled.wait(timeout=timeout)


class WorkerHealthReport:
    """Health report for an async worker.

    Tracks the worker's activity and status for diagnostics.
    """

    def __init__(self, worker_id: str) -> None:
        if not worker_id or not worker_id.strip():
            raise ValueError("worker_id must be a non-empty string")
        self.worker_id: str = worker_id
        self.last_heartbeat: datetime = datetime.now(UTC)
        self.jobs_processed: int = 0
        self.jobs_failed: int = 0
        self.jobs_cancelled: int = 0
        self.is_running: bool = False
        self.started_at: datetime = datetime.now(UTC)

    def record_heartbeat(self) -> None:
        """Record a heartbeat."""
        self.last_heartbeat = datetime.now(UTC)

    def record_job_completed(self) -> None:
        """Record a successfully completed job."""
        self.jobs_processed += 1
        self.record_heartbeat()

    def record_job_failed(self) -> None:
        """Record a failed job."""
        self.jobs_failed += 1
        self.record_heartbeat()

    def record_job_cancelled(self) -> None:
        """Record a cancelled job."""
        self.jobs_cancelled += 1
        self.record_heartbeat()

    def as_dict(self) -> dict[str, Any]:
        """Return a dictionary representation for diagnostics."""
        return {
            "worker_id": self.worker_id,
            "last_heartbeat": self.last_heartbeat.isoformat(),
            "jobs_processed": self.jobs_processed,
            "jobs_failed": self.jobs_failed,
            "jobs_cancelled": self.jobs_cancelled,
            "is_running": self.is_running,
            "started_at": self.started_at.isoformat(),
        }


class AsyncWorkerConfig:
    """Configuration for the async worker service."""

    def __init__(
        self,
        enabled: bool = False,
        max_workers: int = 4,
        poll_interval: float = 1.0,
        job_timeout: int = 3600,
        job_ttl: int = 86400,
    ) -> None:
        if not isinstance(enabled, bool):
            raise TypeError(f"enabled must be bool, got {type(enabled)}")
        if not isinstance(max_workers, int) or max_workers < 1:
            raise ValueError("max_workers must be a positive integer")
        if not isinstance(poll_interval, (int, float)) or poll_interval <= 0:
            raise ValueError("poll_interval must be a positive number")
        if not isinstance(job_timeout, int) or job_timeout < 1:
            raise ValueError("job_timeout must be a positive integer")
        if not isinstance(job_ttl, int) or job_ttl < 1:
            raise ValueError("job_ttl must be a positive integer")
        self.enabled: bool = enabled
        self.max_workers: int = max_workers
        self.poll_interval: float = float(poll_interval)
        self.job_timeout: int = job_timeout
        self.job_ttl: int = job_ttl


class InMemoryJobStore:
    """Thread-safe in-memory job store for async jobs.

    Provides CRUD operations for jobs, used by the worker and tests.
    Production implementations would back this with the database.
    """

    def __init__(self) -> None:
        self._jobs: dict[str, AsyncJob] = {}
        self._lock = threading.Lock()

    def add(self, job: AsyncJob) -> None:
        """Add a job to the store.

        Args:
            job: The job to add.

        Raises:
            ValueError: If a job with the same ID already exists.
        """
        if not isinstance(job, AsyncJob):
            raise TypeError(f"job must be AsyncJob, got {type(job)}")
        with self._lock:
            if job.job_id in self._jobs:
                raise ValueError(f"Job {job.job_id} already exists")
            self._jobs[job.job_id] = job

    def get(self, job_id: str) -> AsyncJob | None:
        """Get a job by ID.

        Args:
            job_id: The job ULID.

        Returns:
            The job, or None if not found.
        """
        if not job_id:
            raise ValueError("job_id must be a non-empty string")
        with self._lock:
            return self._jobs.get(job_id)

    def list_by_status(self, status: AsyncJobStatus) -> list[AsyncJob]:
        """List jobs with a specific status.

        Args:
            status: The status to filter by.

        Returns:
            List of jobs with the given status.
        """
        if not isinstance(status, AsyncJobStatus):
            raise TypeError(f"status must be AsyncJobStatus, got {type(status)}")
        with self._lock:
            return [j for j in self._jobs.values() if j.status == status]

    def list_all(self) -> list[AsyncJob]:
        """List all jobs."""
        with self._lock:
            return list(self._jobs.values())

    def update(self, job: AsyncJob) -> None:
        """Update a job in the store.

        Args:
            job: The job with updated fields.

        Raises:
            ValueError: If the job does not exist.
        """
        if not isinstance(job, AsyncJob):
            raise TypeError(f"job must be AsyncJob, got {type(job)}")
        with self._lock:
            if job.job_id not in self._jobs:
                raise ValueError(f"Job {job.job_id} not found")
            self._jobs[job.job_id] = job

    def remove(self, job_id: str) -> None:
        """Remove a job from the store.

        Args:
            job_id: The job ULID.
        """
        if not job_id:
            raise ValueError("job_id must be a non-empty string")
        with self._lock:
            self._jobs.pop(job_id, None)

    def count(self) -> int:
        """Return the total number of jobs."""
        with self._lock:
            return len(self._jobs)

    def snapshot_counts(self) -> dict[str, int]:
        """Return all status counts atomically in a single lock acquisition.

        Returns:
            Dictionary with ``queued``, ``running``, and ``total`` keys.
        """
        with self._lock:
            queued = sum(
                1 for j in self._jobs.values() if j.status == AsyncJobStatus.QUEUED
            )
            running = sum(
                1 for j in self._jobs.values() if j.status == AsyncJobStatus.RUNNING
            )
            return {"queued": queued, "running": running, "total": len(self._jobs)}

    def remove_expired(self, ttl_seconds: int) -> int:
        """Remove terminal jobs whose ``finished_at`` exceeds TTL.

        Performs a single-pass removal under one lock acquisition to
        avoid N+1 locking overhead.

        Args:
            ttl_seconds: Retention period in seconds.

        Returns:
            Number of jobs removed.
        """
        now = datetime.now(UTC)
        ttl_delta = timedelta(seconds=ttl_seconds)
        with self._lock:
            expired_ids = [
                jid
                for jid, job in self._jobs.items()
                if (
                    job.is_terminal
                    and job.finished_at is not None
                    and now - job.finished_at > ttl_delta
                )
            ]
            for jid in expired_ids:
                del self._jobs[jid]
            return len(expired_ids)


class AsyncWorker:
    """Async worker service for background job execution.

    Uses a ``ThreadPoolExecutor`` to run jobs concurrently up to
    ``max_workers``.  The poll loop dispatches queued jobs to the pool
    and respects the concurrency limit.

    **Cancellation contract:** ``cancel_job`` only signals the
    cancellation token.  The ``pickup_and_execute`` method that owns
    the job is responsible for all state transitions, preventing the
    race condition where two threads attempt conflicting transitions
    on the same job object.
    """

    def __init__(
        self,
        config: AsyncWorkerConfig,
        job_store: InMemoryJobStore,
        job_executor: Any | None = None,
    ) -> None:
        if not isinstance(config, AsyncWorkerConfig):
            raise TypeError(f"config must be AsyncWorkerConfig, got {type(config)}")
        if not isinstance(job_store, InMemoryJobStore):
            raise TypeError(
                f"job_store must be InMemoryJobStore, got {type(job_store)}"
            )
        self._config = config
        self._job_store = job_store
        self._job_executor = job_executor
        self._worker_id = f"worker-{uuid4().hex[:8]}"
        self._shutdown_event = threading.Event()
        self._cancellation_tokens: dict[str, CancellationToken] = {}
        self._tokens_lock = threading.Lock()
        self._health = WorkerHealthReport(self._worker_id)
        self._poll_thread: threading.Thread | None = None
        self._thread_pool: ThreadPoolExecutor | None = None
        self._futures: dict[str, Future[None]] = {}
        self._futures_lock = threading.Lock()
        self._original_sigint: Any = None
        self._original_sigterm: Any = None

    @property
    def config(self) -> AsyncWorkerConfig:
        """Return the worker configuration."""
        return self._config

    @property
    def worker_id(self) -> str:
        """Return the worker ID."""
        return self._worker_id

    @property
    def health(self) -> WorkerHealthReport:
        """Return the worker health report."""
        return self._health

    @property
    def is_running(self) -> bool:
        """Check if the worker is running."""
        return self._health.is_running

    def start(self) -> None:
        """Start the worker polling loop.

        Creates a ``ThreadPoolExecutor`` for concurrent job execution,
        installs signal handlers for graceful shutdown, and starts the
        background polling thread.
        """
        if self._health.is_running:
            return
        self._health.is_running = True
        self._shutdown_event.clear()
        self._thread_pool = ThreadPoolExecutor(
            max_workers=self._config.max_workers,
            thread_name_prefix=f"job-{self._worker_id}",
        )
        self._install_signal_handlers()
        self._poll_thread = threading.Thread(
            target=self._poll_loop,
            name=f"async-worker-{self._worker_id}",
            daemon=True,
        )
        self._poll_thread.start()
        logger.info("AsyncWorker %s started", self._worker_id)

    def stop(self) -> None:
        """Gracefully stop the worker.

        Signals the polling loop to stop, cancels all running jobs,
        shuts down the thread pool, and waits for the poll thread to
        finish.
        """
        if not self._health.is_running:
            return
        logger.info("AsyncWorker %s stopping", self._worker_id)
        self._shutdown_event.set()
        self._cancel_all_running_jobs()
        if self._thread_pool is not None:
            self._thread_pool.shutdown(wait=True, cancel_futures=True)
            self._thread_pool = None
        if self._poll_thread is not None:
            self._poll_thread.join(timeout=self._config.poll_interval * 3)
        self._health.is_running = False
        self._restore_signal_handlers()
        logger.info("AsyncWorker %s stopped", self._worker_id)

    def request_shutdown(self) -> None:
        """Request a graceful shutdown (signal-handler safe)."""
        self._shutdown_event.set()

    def pickup_and_execute(self, job: AsyncJob) -> None:
        """Pick up a job and execute it.

        This method owns **all** state transitions for the job.
        External cancellation (via ``cancel_job``) only signals the
        cancellation token; this method checks the token and performs
        the appropriate transition.

        Args:
            job: The job to execute.

        Raises:
            TypeError: If job is not an AsyncJob.
        """
        if not isinstance(job, AsyncJob):
            raise TypeError(f"job must be AsyncJob, got {type(job)}")

        token = CancellationToken()
        with self._tokens_lock:
            self._cancellation_tokens[job.job_id] = token

        try:
            # If the job was already cancelled before we got here, skip.
            if job.is_terminal:
                return

            job.mark_running(self._worker_id)
            self._job_store.update(job)
            logger.info("Job %s picked up by %s", job.job_id, self._worker_id)

            # Execute the job
            if self._job_executor is not None:
                self._job_executor(job, token)

            # Determine final state — guard against concurrent cancellation
            # that may have already transitioned this job via another path.
            if job.is_terminal:
                # Another thread already moved the job to a terminal state
                # (should not happen with the new contract, but defensive).
                logger.info("Job %s already terminal: %s", job.job_id, job.status.value)
            elif token.is_cancelled:
                job.mark_cancelled()
                self._job_store.update(job)
                self._health.record_job_cancelled()
                logger.info("Job %s cancelled", job.job_id)
            else:
                job.mark_succeeded()
                self._job_store.update(job)
                self._health.record_job_completed()
                logger.info("Job %s succeeded", job.job_id)
        except InvalidJobTransitionError:
            logger.warning("Job %s transition conflict (already terminal)", job.job_id)
        except Exception as exc:
            error_msg = redact_value(f"{type(exc).__name__}: {exc}")
            logger.exception("Job %s failed with error", job.job_id)
            try:
                job.mark_failed(error=error_msg)
                self._job_store.update(job)
            except InvalidJobTransitionError:
                pass
            self._health.record_job_failed()
        finally:
            with self._tokens_lock:
                self._cancellation_tokens.pop(job.job_id, None)
            with self._futures_lock:
                self._futures.pop(job.job_id, None)

    def cancel_job(self, job_id: str) -> bool:
        """Cancel a specific job.

        Only signals the cancellation token.  The ``pickup_and_execute``
        method that owns the job performs the actual state transition,
        preventing a race where two threads attempt conflicting
        transitions on the same job object.

        For jobs that are still ``queued`` (not yet picked up), this
        method directly transitions them to ``cancelled`` since no
        worker thread owns them yet.

        Args:
            job_id: The job ULID to cancel.

        Returns:
            True if the job was cancelled or signalled, False otherwise.
        """
        if not job_id:
            raise ValueError("job_id must be a non-empty string")

        # Signal the cancellation token for running jobs
        with self._tokens_lock:
            token = self._cancellation_tokens.get(job_id)
            if token is not None:
                token.cancel()
                # Job is owned by a worker thread — let it handle the
                # state transition via pickup_and_execute.
                return True

        # For queued jobs not yet picked up, cancel directly.
        job = self._job_store.get(job_id)
        if job is not None and not job.is_terminal:
            try:
                job.mark_cancelled()
                self._job_store.update(job)
                return True
            except InvalidJobTransitionError:
                pass
        return False

    def detect_stuck_jobs(self) -> list[AsyncJob]:
        """Detect and fail jobs that have exceeded the heartbeat TTL.

        A job is considered stuck if it is in RUNNING state and its
        last heartbeat is older than ``job_timeout`` seconds.

        Returns:
            List of jobs that were marked as failed.
        """
        running_jobs = self._job_store.list_by_status(AsyncJobStatus.RUNNING)
        stuck: list[AsyncJob] = []
        now = datetime.now(UTC)
        timeout_delta = timedelta(seconds=self._config.job_timeout)

        for job in running_jobs:
            heartbeat = job.last_heartbeat or job.started_at or job.created_at
            if now - heartbeat > timeout_delta:
                try:
                    job.mark_failed(error="Heartbeat timeout — job stuck")
                    self._job_store.update(job)
                    stuck.append(job)
                    logger.warning(
                        "Job %s marked failed: heartbeat timeout", job.job_id
                    )
                except InvalidJobTransitionError:
                    pass
        return stuck

    def cleanup_completed_jobs(self) -> int:
        """Remove completed jobs older than the retention TTL.

        Delegates to ``InMemoryJobStore.remove_expired`` for a
        single-pass, single-lock removal.

        Returns:
            Number of jobs cleaned up.
        """
        cleaned = self._job_store.remove_expired(self._config.job_ttl)
        if cleaned > 0:
            logger.debug("Cleaned up %d expired job(s)", cleaned)
        return cleaned

    def get_health_report(self) -> dict[str, Any]:
        """Return a health report for diagnostics.

        Uses ``snapshot_counts`` on the job store to obtain all queue
        depth metrics under a single lock acquisition for consistency.

        Returns:
            Dictionary with worker health information.
        """
        self._health.record_heartbeat()
        report = self._health.as_dict()
        report["config"] = {
            "enabled": self._config.enabled,
            "max_workers": self._config.max_workers,
            "poll_interval": self._config.poll_interval,
            "job_timeout": self._config.job_timeout,
            "job_ttl": self._config.job_ttl,
        }
        counts = self._job_store.snapshot_counts()
        report["queued_jobs"] = counts["queued"]
        report["running_jobs"] = counts["running"]
        report["total_jobs"] = counts["total"]
        return report

    # -- Private methods ----------------------------------------------------

    def _poll_loop(self) -> None:
        """Main polling loop that dispatches queued jobs to the thread pool.

        Each job is submitted to the ``ThreadPoolExecutor`` for truly
        concurrent execution up to ``max_workers``.
        """
        while not self._shutdown_event.is_set():
            try:
                self._health.record_heartbeat()
                # Determine available capacity in the pool
                with self._futures_lock:
                    active = sum(1 for f in self._futures.values() if not f.done())
                capacity = self._config.max_workers - active
                if capacity > 0:
                    queued = self._job_store.list_by_status(AsyncJobStatus.QUEUED)
                    for job in queued[:capacity]:
                        if self._shutdown_event.is_set():
                            break
                        self._dispatch_job(job)
            except Exception:
                logger.exception("Error in poll loop")
            self._shutdown_event.wait(timeout=self._config.poll_interval)

    def _dispatch_job(self, job: AsyncJob) -> None:
        """Submit a job to the thread pool for concurrent execution."""
        if self._thread_pool is None:
            # Fallback for direct pickup_and_execute calls (e.g., tests)
            self.pickup_and_execute(job)
            return
        with self._futures_lock:
            if job.job_id in self._futures:
                return  # Already dispatched
            future: Future[None] = self._thread_pool.submit(
                self.pickup_and_execute, job
            )
            self._futures[job.job_id] = future

    def _cancel_all_running_jobs(self) -> None:
        """Cancel all currently running jobs."""
        with self._tokens_lock:
            for token in self._cancellation_tokens.values():
                token.cancel()

    def _signal_handler(self, signum: int, _frame: Any) -> None:
        """Handle shutdown signals."""
        logger.info("Received signal %d, requesting shutdown", signum)
        self.request_shutdown()

    def _install_signal_handlers(self) -> None:
        """Install signal handlers for graceful shutdown."""
        try:
            self._original_sigint = signal.getsignal(signal.SIGINT)
            self._original_sigterm = signal.getsignal(signal.SIGTERM)
            signal.signal(signal.SIGINT, self._signal_handler)
            signal.signal(signal.SIGTERM, self._signal_handler)
        except (OSError, ValueError):
            # Cannot set signal handlers from non-main thread
            pass

    def _restore_signal_handlers(self) -> None:
        """Restore original signal handlers."""
        try:
            if self._original_sigint is not None:
                signal.signal(signal.SIGINT, self._original_sigint)
            if self._original_sigterm is not None:
                signal.signal(signal.SIGTERM, self._original_sigterm)
        except (OSError, ValueError):
            pass
