"""Job cancellation and timeout enforcement module for Oclawma.

This module provides:
- CancellationToken: Cooperative cancellation mechanism
- TimeoutManager: Hard and soft timeout handling with force kill
"""

from __future__ import annotations

import logging
import os
import signal
import threading
import time
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass
from enum import Enum
from typing import Any, Callable

logger = logging.getLogger(__name__)


class CancellationReason(str, Enum):
    """Reason for job cancellation."""

    USER_REQUEST = "user_request"
    TIMEOUT_SOFT = "timeout_soft"
    TIMEOUT_HARD = "timeout_hard"
    SHUTDOWN = "shutdown"
    PREEMPTION = "preemption"


class CancellationError(Exception):
    """Raised when a job is cancelled."""

    def __init__(self, reason: CancellationReason, message: str = "") -> None:
        self.reason = reason
        self.message = message
        super().__init__(f"Job cancelled: {reason.value} - {message}")


class CancellationToken:
    """Cooperative cancellation token for jobs.

    Jobs can check this token periodically to see if cancellation
    has been requested and exit gracefully.

    Example:
        >>> token = CancellationToken()
        >>> def long_running_job(token: CancellationToken):
        ...     for i in range(100):
        ...         token.check_cancelled()  # Raises if cancelled
        ...         # or
        ...         if token.is_cancelled():
        ...             break
        ...         do_work(i)
        >>> # From another thread:
        >>> token.request_cancel(CancellationReason.USER_REQUEST)
    """

    def __init__(self) -> None:
        self._cancelled = False
        self._reason: CancellationReason | None = None
        self._message: str = ""
        self._cancelled_at: float | None = None
        self._lock = threading.Lock()
        self._callbacks: list[Callable[[CancellationReason, str], None]] = []

    def request_cancel(
        self,
        reason: CancellationReason = CancellationReason.USER_REQUEST,
        message: str = "",
    ) -> bool:
        """Request cancellation of the job.

        Args:
            reason: The reason for cancellation
            message: Optional message explaining the cancellation

        Returns:
            True if cancellation was newly requested, False if already cancelled
        """
        with self._lock:
            if self._cancelled:
                return False

            self._cancelled = True
            self._reason = reason
            self._message = message
            self._cancelled_at = time.time()

            # Notify callbacks
            callbacks = self._callbacks.copy()

        # Call callbacks outside the lock to avoid deadlocks
        for callback in callbacks:
            try:
                callback(reason, message)
            except Exception as e:
                logger.warning(f"Cancellation callback failed: {e}")

        logger.debug(f"Cancellation requested: {reason.value} - {message}")
        return True

    def is_cancelled(self) -> bool:
        """Check if cancellation has been requested.

        Returns:
            True if cancellation was requested
        """
        with self._lock:
            return self._cancelled

    def check_cancelled(self) -> None:
        """Check if cancelled and raise CancellationError if so.

        This is a convenience method for cooperative cancellation.
        Jobs should call this periodically to respond to cancellation.

        Raises:
            CancellationError: If cancellation has been requested
        """
        with self._lock:
            if self._cancelled:
                raise CancellationError(
                    self._reason or CancellationReason.USER_REQUEST, self._message
                )

    @property
    def reason(self) -> CancellationReason | None:
        """Get the cancellation reason."""
        with self._lock:
            return self._reason

    @property
    def message(self) -> str:
        """Get the cancellation message."""
        with self._lock:
            return self._message

    @property
    def cancelled_at(self) -> float | None:
        """Get the timestamp when cancellation was requested."""
        with self._lock:
            return self._cancelled_at

    def add_callback(self, callback: Callable[[CancellationReason, str], None]) -> None:
        """Add a callback to be invoked when cancellation is requested.

        Args:
            callback: Function to call with (reason, message) when cancelled
        """
        with self._lock:
            self._callbacks.append(callback)

        # If already cancelled, invoke immediately
        if self._cancelled:
            try:
                callback(self._reason or CancellationReason.USER_REQUEST, self._message)
            except Exception as e:
                logger.warning(f"Cancellation callback failed: {e}")

    def remove_callback(self, callback: Callable[[CancellationReason, str], None]) -> bool:
        """Remove a cancellation callback.

        Args:
            callback: The callback to remove

        Returns:
            True if callback was found and removed
        """
        with self._lock:
            if callback in self._callbacks:
                self._callbacks.remove(callback)
                return True
            return False

    def reset(self) -> None:
        """Reset the cancellation state.

        This is useful for reusing tokens in retry scenarios.
        """
        with self._lock:
            self._cancelled = False
            self._reason = None
            self._message = ""
            self._cancelled_at = None
            self._callbacks.clear()


@dataclass
class TimeoutConfig:
    """Configuration for timeout handling."""

    soft_timeout: float | None = None
    """Soft timeout in seconds. Job is notified but can continue."""

    hard_timeout: float | None = None
    """Hard timeout in seconds. Job is forcefully terminated."""

    grace_period: float = 5.0
    """Seconds to wait between soft and hard timeout."""

    def __post_init__(self) -> None:
        if self.soft_timeout is not None and self.soft_timeout <= 0:
            raise ValueError("soft_timeout must be positive")
        if self.hard_timeout is not None and self.hard_timeout <= 0:
            raise ValueError("hard_timeout must be positive")
        if self.grace_period < 0:
            raise ValueError("grace_period must be non-negative")


class TimeoutManager:
    """Manages job timeouts with soft and hard timeout handling.

    Features:
    - Soft timeout: Notifies job via cancellation token
    - Hard timeout: Forcefully terminates job (SIGTERM on Unix)
    - Force kill: SIGKILL if grace period expires after hard timeout

    Example:
        >>> config = TimeoutConfig(soft_timeout=30, hard_timeout=60, grace_period=10)
        >>> manager = TimeoutManager(config)
        >>> token = manager.start_monitoring(job_id=1)
        >>> # Job should check token periodically
        >>> try:
        ...     job.run(token)
        ... except CancellationError:
        ...     print("Job was cancelled")
        >>> manager.stop_monitoring()
    """

    def __init__(self, config: TimeoutConfig | None = None) -> None:
        self.config = config or TimeoutConfig()
        self._token = CancellationToken()
        self._timer: threading.Timer | None = None
        self._hard_timer: threading.Timer | None = None
        self._force_kill_timer: threading.Timer | None = None
        self._lock = threading.Lock()
        self._job_id: int | str | None = None
        self._start_time: float | None = None
        self._target_pid: int | None = None
        self._executor: ThreadPoolExecutor | None = None
        self._future: Future[Any] | None = None
        self._on_hard_timeout: Callable[[], None] | None = None
        self._on_force_kill: Callable[[], None] | None = None

    @property
    def token(self) -> CancellationToken:
        """Get the cancellation token."""
        return self._token

    def set_timeout(
        self,
        soft_timeout: float | None = None,
        hard_timeout: float | None = None,
        grace_period: float | None = None,
    ) -> None:
        """Update timeout configuration.

        Args:
            soft_timeout: New soft timeout in seconds
            hard_timeout: New hard timeout in seconds
            grace_period: New grace period in seconds
        """
        with self._lock:
            if soft_timeout is not None:
                self.config.soft_timeout = soft_timeout
            if hard_timeout is not None:
                self.config.hard_timeout = hard_timeout
            if grace_period is not None:
                self.config.grace_period = grace_period

    def cancel_after(
        self, seconds: float, reason: CancellationReason = CancellationReason.TIMEOUT_SOFT
    ) -> None:
        """Request cancellation after a specified time.

        This is a convenience method for delayed cancellation.

        Args:
            seconds: Time in seconds until cancellation
            reason: Reason for cancellation
        """

        def delayed_cancel():
            self._token.request_cancel(reason, f"Cancelled after {seconds}s")

        timer = threading.Timer(seconds, delayed_cancel)
        timer.daemon = True
        timer.start()

    def extend_timeout(self, additional_seconds: float) -> bool:
        """Extend the current timeout by additional seconds.

        This restarts the monitoring timers with the new timeout values.

        Args:
            additional_seconds: Seconds to add to current timeouts

        Returns:
            True if timeout was extended, False if already expired
        """
        with self._lock:
            if self._token.is_cancelled():
                return False

            # Calculate remaining time and add extension
            elapsed = time.time() - (self._start_time or time.time())

            if self.config.soft_timeout is not None:
                remaining_soft = max(0, self.config.soft_timeout - elapsed)
                self.config.soft_timeout = remaining_soft + additional_seconds

            if self.config.hard_timeout is not None:
                remaining_hard = max(0, self.config.hard_timeout - elapsed)
                self.config.hard_timeout = remaining_hard + additional_seconds

            # Restart monitoring with new timeouts
            self._stop_timers()
            self._start_timers()

            logger.debug(f"Timeout extended by {additional_seconds}s")
            return True

    def start_monitoring(
        self,
        job_id: int | str | None = None,
        target_pid: int | None = None,
        on_hard_timeout: Callable[[], None] | None = None,
        on_force_kill: Callable[[], None] | None = None,
    ) -> CancellationToken:
        """Start monitoring timeouts for a job.

        Args:
            job_id: Optional job identifier for logging
            target_pid: Process ID to signal on hard timeout (Unix only)
            on_hard_timeout: Callback when hard timeout is reached
            on_force_kill: Callback when force kill is triggered

        Returns:
            CancellationToken for the job to check
        """
        with self._lock:
            self._job_id = job_id
            self._start_time = time.time()
            self._target_pid = target_pid
            self._on_hard_timeout = on_hard_timeout
            self._on_force_kill = on_force_kill

            self._start_timers()

            logger.debug(f"Started timeout monitoring for job {job_id}")
            return self._token

    def _start_timers(self) -> None:
        """Start the timeout monitoring timers."""
        # Soft timeout timer
        if self.config.soft_timeout is not None:
            self._timer = threading.Timer(
                self.config.soft_timeout,
                self._on_soft_timeout,
            )
            self._timer.daemon = True
            self._timer.start()

        # Hard timeout timer
        if self.config.hard_timeout is not None:
            self._hard_timer = threading.Timer(
                self.config.hard_timeout,
                self._on_hard_timeout_triggered,
            )
            self._hard_timer.daemon = True
            self._hard_timer.start()

    def stop_monitoring(self) -> None:
        """Stop timeout monitoring and cancel all timers."""
        with self._lock:
            self._stop_timers()
            logger.debug(f"Stopped timeout monitoring for job {self._job_id}")

    def _stop_timers(self) -> None:
        """Cancel all active timers."""
        for timer in [self._timer, self._hard_timer, self._force_kill_timer]:
            if timer is not None:
                timer.cancel()
        self._timer = None
        self._hard_timer = None
        self._force_kill_timer = None

    def _on_soft_timeout(self) -> None:
        """Handle soft timeout - request cooperative cancellation."""
        logger.warning(
            f"Soft timeout reached for job {self._job_id}. " f"Requesting graceful cancellation."
        )
        self._token.request_cancel(
            CancellationReason.TIMEOUT_SOFT,
            f"Soft timeout ({self.config.soft_timeout}s) reached",
        )

    def _on_hard_timeout_triggered(self) -> None:
        """Handle hard timeout - forceful termination."""
        logger.error(
            f"Hard timeout reached for job {self._job_id}. " f"Initiating forceful termination."
        )
        self._token.request_cancel(
            CancellationReason.TIMEOUT_HARD,
            f"Hard timeout ({self.config.hard_timeout}s) reached",
        )

        # Signal the target process if specified
        if self._target_pid is not None:
            self._signal_process(self._target_pid, signal.SIGTERM)

        # Call the hard timeout callback
        if self._on_hard_timeout:
            try:
                self._on_hard_timeout()
            except Exception as e:
                logger.error(f"Hard timeout callback failed: {e}")

        # Schedule force kill after grace period
        if self.config.grace_period > 0:
            self._force_kill_timer = threading.Timer(
                self.config.grace_period,
                self._on_force_kill_triggered,
            )
            self._force_kill_timer.daemon = True
            self._force_kill_timer.start()

    def _on_force_kill_triggered(self) -> None:
        """Handle force kill - SIGKILL the process."""
        logger.critical(
            f"Force kill triggered for job {self._job_id} "
            f"after grace period of {self.config.grace_period}s"
        )

        # SIGKILL the target process
        if self._target_pid is not None:
            self._signal_process(self._target_pid, signal.SIGKILL)

        # Call the force kill callback
        if self._on_force_kill:
            try:
                self._on_force_kill()
            except Exception as e:
                logger.error(f"Force kill callback failed: {e}")

    def _signal_process(self, pid: int, sig: int) -> bool:
        """Send a signal to a process.

        Args:
            pid: Process ID
            sig: Signal number

        Returns:
            True if signal was sent successfully
        """
        try:
            os.kill(pid, sig)
            logger.debug(f"Sent signal {sig} to process {pid}")
            return True
        except ProcessLookupError:
            logger.warning(f"Process {pid} not found")
            return False
        except PermissionError:
            logger.error(f"Permission denied signaling process {pid}")
            return False
        except Exception as e:
            logger.error(f"Failed to signal process {pid}: {e}")
            return False

    def is_expired(self) -> bool:
        """Check if the timeout has expired.

        Returns:
            True if hard timeout has been reached
        """
        with self._lock:
            if self._start_time is None or self.config.hard_timeout is None:
                return False
            elapsed = time.time() - self._start_time
            return elapsed >= self.config.hard_timeout

    def get_remaining_time(self) -> float | None:
        """Get remaining time before hard timeout.

        Returns:
            Seconds remaining, or None if no hard timeout
        """
        with self._lock:
            if self._start_time is None or self.config.hard_timeout is None:
                return None
            elapsed = time.time() - self._start_time
            return max(0, self.config.hard_timeout - elapsed)

    def get_elapsed_time(self) -> float:
        """Get elapsed time since monitoring started.

        Returns:
            Seconds elapsed
        """
        with self._lock:
            if self._start_time is None:
                return 0.0
            return time.time() - self._start_time


def run_with_timeout(
    func: Callable[[CancellationToken], Any],
    timeout_config: TimeoutConfig,
    *args: Any,
    **kwargs: Any,
) -> Any:
    """Run a function with timeout enforcement.

    This is a convenience function that runs a function in a thread
    with timeout monitoring.

    Args:
        func: Function to run, should accept CancellationToken as first arg
        timeout_config: Timeout configuration
        *args: Additional positional arguments for func
        **kwargs: Additional keyword arguments for func

    Returns:
        The return value of func

    Raises:
        CancellationError: If the job is cancelled due to timeout
        Exception: Any exception raised by func
    """
    manager = TimeoutManager(timeout_config)
    token = manager.start_monitoring()

    result: Any = None
    error: Exception | None = None

    def wrapper():
        nonlocal result, error
        try:
            result = func(token, *args, **kwargs)
        except Exception as e:
            error = e

    thread = threading.Thread(target=wrapper)
    thread.start()
    thread.join(timeout_config.hard_timeout)

    if thread.is_alive():
        # Hard timeout reached but thread still running
        manager.stop_monitoring()
        raise CancellationError(
            CancellationReason.TIMEOUT_HARD,
            f"Job exceeded hard timeout of {timeout_config.hard_timeout}s",
        )

    manager.stop_monitoring()

    if error is not None:
        raise error

    # Check if cancelled
    token.check_cancelled()

    return result
