from __future__ import annotations

import logging
import os
import signal
import socket
import time
import traceback
from collections.abc import Iterable
from datetime import datetime
from enum import Enum
from multiprocessing import Process, get_context
from multiprocessing.process import BaseProcess
from uuid import uuid4

from redis import ConnectionPool, Redis
from redis.client import Pipeline

from .connections import parse_connection
from .defaults import DEFAULT_LOGGING_DATE_FORMAT, DEFAULT_LOGGING_FORMAT, DEFAULT_SCHEDULER_FALLBACK_PERIOD
from .exceptions import SchedulerNotFound
from .job import Job
from .logutils import setup_loghandlers
from .queue import Queue
from .registry import ScheduledJobRegistry
from .scripts import ACQUIRE_OR_REFRESH_LOCK_SCRIPT, acquire_or_refresh_lock, release_lock
from .serializers import resolve_serializer
from .utils import current_timestamp, decode_redis_hash, now, parse_names, utcformat, utcparse

ForkProcess: type[BaseProcess]
try:
    ForkProcess = get_context('fork').Process
except ValueError:
    ForkProcess = Process

SCHEDULER_KEY_TEMPLATE = 'rq:scheduler:%s'
SCHEDULER_LOCKING_KEY_TEMPLATE = 'rq:scheduler-lock:%s'


class SchedulerStatus(str, Enum):
    STARTED = 'started'
    WORKING = 'working'
    STOPPED = 'stopped'


class RQScheduler:
    # STARTED: scheduler has been started but sleeping
    # WORKING: scheduler is in the midst of scheduling jobs
    # STOPPED: scheduler is in stopped condition

    Status = SchedulerStatus

    def __init__(
        self,
        queues,
        connection: Redis,
        interval=1,
        logging_level: str | int = logging.INFO,
        date_format=DEFAULT_LOGGING_DATE_FORMAT,
        log_format=DEFAULT_LOGGING_FORMAT,
        serializer=None,
        name: str | None = None,
    ):
        self._queue_names = set(parse_names(queues))
        self._acquired_locks: set[str] = set()
        self._scheduled_job_registries: list[ScheduledJobRegistry] = []
        self.lock_acquisition_time = None
        self._connection_class, self._pool_class, self._pool_kwargs = parse_connection(connection)
        self.serializer = resolve_serializer(serializer)

        # Identity, stable across the fork (name is pickled to the child process).
        self.name: str = name or uuid4().hex
        self.hostname: str = socket.gethostname()
        self.created_at: datetime = now()
        self.pid: int = 0  # set in register_birth(), once the scheduler process has forked
        self.last_heartbeat: datetime | None = None

        self._connection = None
        self.interval = interval
        self._stop_requested = False
        self._status = self.Status.STOPPED
        self._process = None
        self.log = logging.getLogger(__name__)
        setup_loghandlers(
            level=logging_level,
            name=__name__,
            log_format=log_format,
            date_format=date_format,
        )

    @property
    def connection(self):
        if self._connection:
            return self._connection
        self._connection = self._connection_class(
            connection_pool=ConnectionPool(connection_class=self._pool_class, **self._pool_kwargs)
        )
        return self._connection

    @property
    def acquired_locks(self):
        return self._acquired_locks

    @property
    def status(self):
        return self._status

    @property
    def key(self) -> str:
        """Redis key holding this scheduler's metadata hash."""
        return SCHEDULER_KEY_TEMPLATE % self.name

    @property
    def should_reacquire_locks(self):
        """Returns True if lock_acquisition_time is longer than 10 minutes ago"""
        if self._queue_names == self.acquired_locks:
            return False
        if not self.lock_acquisition_time:
            return True
        return (datetime.now() - self.lock_acquisition_time).total_seconds() > DEFAULT_SCHEDULER_FALLBACK_PERIOD

    def acquire_locks(self, auto_start=False):
        """Acquire or refresh scheduler locks, returning queue names verified as owned
        (newly acquired or reclaimed)."""
        successful_locks = set()
        self.log.debug('Acquiring scheduler lock for %s', ', '.join(self._queue_names))
        for name in self._queue_names:
            outcome = acquire_or_refresh_lock(
                self.connection, self.get_locking_key(name), self.name, self.interval + 60
            )
            if outcome == 'acquired':
                self.log.info('Acquired scheduler lock for %s', name)
                successful_locks.add(name)
            elif outcome == 'refreshed':
                self.log.debug('Refreshed scheduler lock for %s', name)
                successful_locks.add(name)

        # Always reset _scheduled_job_registries when acquiring locks
        self._scheduled_job_registries = []
        self._acquired_locks = successful_locks
        self.lock_acquisition_time = datetime.now()

        # If auto_start is requested and scheduler is not started,
        # run self.start()
        if self._acquired_locks and auto_start:
            if not self._process or not self._process.is_alive():
                self.start()

        return successful_locks

    def prepare_registries(self, queue_names: Iterable[str] | None = None):
        """Prepare scheduled job registries for use"""
        self._scheduled_job_registries = []
        if not queue_names:
            queue_names = self._acquired_locks
        for name in queue_names:
            self._scheduled_job_registries.append(
                ScheduledJobRegistry(name, connection=self.connection, serializer=self.serializer)
            )

    @classmethod
    def get_locking_key(cls, name: str):
        """Returns scheduler key for a given queue name"""
        return SCHEDULER_LOCKING_KEY_TEMPLATE % name

    def to_dict(self) -> dict:
        """Serialize this scheduler's metadata for storage in its Redis hash."""
        assert self.last_heartbeat is not None
        return {
            'name': self.name,
            'hostname': self.hostname,
            'pid': str(self.pid),
            'queues': ','.join(self._queue_names),
            'created_at': utcformat(self.created_at),
            'last_heartbeat': utcformat(self.last_heartbeat),
        }

    def restore(self, raw_data: dict) -> None:
        """Restore this scheduler's metadata from its Redis hash."""
        obj = decode_redis_hash(raw_data, decode_values=True)
        self.name = obj['name']
        self.hostname = obj['hostname']
        self.pid = int(obj['pid'])
        self._queue_names = set(obj['queues'].split(',')) if obj.get('queues') else set()
        self.created_at = utcparse(obj['created_at'])
        self.last_heartbeat = utcparse(obj['last_heartbeat']) if obj.get('last_heartbeat') else None

    def save(self, pipeline: Pipeline | None = None) -> None:
        """Save this scheduler's metadata hash with a TTL."""
        connection = pipeline if pipeline is not None else self.connection.pipeline()
        connection.hset(self.key, mapping=self.to_dict())
        connection.expire(self.key, self.interval + 60)

        if pipeline is None:
            connection.execute()

    def register_birth(self) -> None:
        """Register this scheduler's birth by writing its metadata hash.

        Idempotent: re-registering the same name (e.g. when the worker restarts a crashed
        scheduler process) overwrites the existing hash rather than erroring.
        """
        self.log.debug('Scheduler %s: registering birth', self.name)
        self.pid = os.getpid()
        self.last_heartbeat = now()
        self.save()

    def register_death(self) -> bool:
        """Register this scheduler's death by deleting its metadata hash.

        Returns:
            True if the scheduler metadata existed and was deleted, False if it was already absent.
        """
        self.log.debug('Scheduler %s: registering death', self.name)
        return bool(self.connection.delete(self.key))

    @classmethod
    def fetch(cls, name: str, connection: Redis) -> RQScheduler:
        """Fetch a scheduler by name, restoring it from its Redis hash."""
        raw_data = connection.hgetall(SCHEDULER_KEY_TEMPLATE % name)
        if not raw_data:
            raise SchedulerNotFound(f"Scheduler with name '{name}' not found")
        scheduler = cls([], connection=connection, name=name)
        scheduler.restore(raw_data)
        return scheduler

    def enqueue_scheduled_jobs(self):
        """Enqueue jobs whose timestamp is in the past"""
        self._status = self.Status.WORKING

        if not self._scheduled_job_registries and self._acquired_locks:
            self.prepare_registries()

        for registry in self._scheduled_job_registries:
            timestamp = current_timestamp()

            # TODO: try to use Lua script to make get_jobs_to_schedule()
            # and remove_jobs() atomic
            job_ids = registry.get_jobs_to_schedule(timestamp)

            if not job_ids:
                continue

            queue = Queue(registry.name, connection=self.connection, serializer=self.serializer)

            with self.connection.pipeline() as pipeline:
                jobs = Job.fetch_many(job_ids, connection=self.connection, serializer=self.serializer)
                for job in jobs:
                    if job is not None:
                        queue._enqueue_job(job, pipeline=pipeline, at_front=job.should_enqueue_at_front())
                for job_id in job_ids:
                    registry.remove(job_id, pipeline=pipeline)
                pipeline.execute()
        self._status = self.Status.STARTED

    def _install_signal_handlers(self):
        """Installs signal handlers for handling SIGINT and SIGTERM
        gracefully.
        """
        signal.signal(signal.SIGINT, self.request_stop)
        signal.signal(signal.SIGTERM, self.request_stop)

    def request_stop(self, signum=None, frame=None):
        """Toggle self._stop_requested that's checked on every loop"""
        self._stop_requested = True

    def heartbeat(self):
        """Refresh the TTL on the scheduler's metadata hash and the locks it owns.
        An expired lock is re-acquired; a lock taken over by another scheduler is
        dropped from `_acquired_locks` so its queue is no longer scheduled."""
        self.log.debug('Scheduler sending heartbeat to %s', ', '.join(self.acquired_locks))
        self.last_heartbeat = now()
        lock_names = sorted(self._acquired_locks)
        lock_script = self.connection.register_script(ACQUIRE_OR_REFRESH_LOCK_SCRIPT)
        with self.connection.pipeline() as pipeline:
            pipeline.hset(self.key, 'last_heartbeat', utcformat(self.last_heartbeat))
            pipeline.expire(self.key, self.interval + 60)
            for name in lock_names:
                lock_script(keys=[self.get_locking_key(name)], args=[self.name, self.interval + 60], client=pipeline)
            results = pipeline.execute()

        lost_locks = set()
        for name, outcome in zip(lock_names, results[2:]):
            if outcome == 0:
                lost_locks.add(name)
            elif outcome == 1:
                # In heartbeat context, a fresh acquire means our lease had expired unclaimed
                self.log.warning('Scheduler lock for %s had expired; re-acquired', name)
        if lost_locks:
            self.log.warning(
                'Scheduler locks for %s were taken over by another scheduler', ', '.join(sorted(lost_locks))
            )
            self._acquired_locks -= lost_locks
            self._scheduled_job_registries = []

    def stop(self):
        self.log.info('Scheduler stopping, releasing locks for %s...', ', '.join(self._acquired_locks))
        self.release_locks()
        self._status = self.Status.STOPPED
        self.register_death()

    def release_locks(self):
        """Release locks still owned by this scheduler, leaving locks held by other schedulers untouched"""
        for name in self._acquired_locks:
            release_lock(self.connection, self.get_locking_key(name), self.name)
        self._acquired_locks = set()

    def start(self):
        self._status = self.Status.STARTED
        # Redis instance can't be pickled across processes so we need to
        # clean this up before forking
        self._connection = None
        self._process = ForkProcess(target=run, args=(self,), name='Scheduler')
        self._process.start()
        return self._process

    def work(self):
        self._install_signal_handlers()
        self.register_birth()

        while True:
            if self._stop_requested:
                self.stop()
                break

            if self.should_reacquire_locks:
                self.acquire_locks()

            self.heartbeat()
            self.enqueue_scheduled_jobs()
            time.sleep(self.interval)


def run(scheduler):
    scheduler.log.info('Scheduler for %s started with PID %s', ', '.join(scheduler._queue_names), os.getpid())
    try:
        scheduler.work()
    except:  # noqa
        scheduler.log.error('Scheduler [PID %s] raised an exception.\n%s', os.getpid(), traceback.format_exc())
        raise
    scheduler.log.info('Scheduler with PID %d has stopped', os.getpid())
