import json
import logging
import time
from datetime import timedelta

from celery.beat import ScheduleEntry, Scheduler
from django.db import DatabaseError, close_old_connections

from task_scheduler.models import PeriodicTask


logger = logging.getLogger(__name__)


class DatabaseScheduler(Scheduler):
    Entry = ScheduleEntry

    def __init__(self, *args, **kwargs):
        self._dirty = {}
        self._modelIds = {}
        self._snapshot = None
        self._nextRefreshAt = 0
        super().__init__(*args, **kwargs)

    def setup_schedule(self):
        schedule, model_ids, snapshot = self._load_schedule()
        self.schedule = schedule
        self._modelIds = model_ids
        self._snapshot = snapshot
        self._nextRefreshAt = time.monotonic() + self.max_interval

    def reserve(self, entry):
        new_entry = super().reserve(entry)
        model_id = self._modelIds.get(entry.name)
        if model_id is not None:
            self._dirty[model_id] = (
                new_entry.last_run_at,
                new_entry.total_run_count,
            )
        return new_entry

    def sync(self):
        if not self._dirty:
            return True

        close_old_connections()
        pending = self._dirty.copy()
        try:
            models = PeriodicTask.objects.in_bulk(pending)
            changed_models = []
            for model_id, (last_run_at, total_run_count) in pending.items():
                model = models.get(model_id)
                if model is None:
                    continue
                model.lastRunAt = last_run_at
                model.totalRunCount = total_run_count
                changed_models.append(model)

            if changed_models:
                PeriodicTask.objects.bulk_update(
                    changed_models,
                    ("lastRunAt", "totalRunCount"),
                )
        except DatabaseError:
            logger.exception("Unable to persist the periodic task run state")
            return False

        for model_id, state in pending.items():
            if self._dirty.get(model_id) == state:
                self._dirty.pop(model_id, None)
        return True

    def tick(self, *args, **kwargs):
        if time.monotonic() >= self._nextRefreshAt:
            self._refresh_schedule()
        return super().tick(*args, **kwargs)

    def _refresh_schedule(self):
        try:
            if not self.sync():
                return
            schedule, model_ids, snapshot = self._load_schedule()
            if snapshot != self._snapshot:
                self.schedule = schedule
                self._modelIds = model_ids
                self._snapshot = snapshot
                self._heap = None
        except DatabaseError:
            logger.exception("Unable to refresh the periodic task schedule")
        finally:
            self._nextRefreshAt = time.monotonic() + self.max_interval

    def _load_schedule(self):
        close_old_connections()
        schedule = {}
        model_ids = {}
        snapshot = []

        for model in PeriodicTask.objects.all().iterator():
            snapshot.append(self._model_snapshot(model))
            if not model.enabled:
                continue
            try:
                entry = self._entry_from_model(model)
            except (TypeError, ValueError) as error:
                logger.error(
                    "Skipping invalid periodic task %s: %s",
                    model.name,
                    error,
                )
                continue
            schedule[model.name] = entry
            model_ids[model.name] = model.pk

        return schedule, model_ids, tuple(snapshot)

    def _entry_from_model(self, model):
        if not isinstance(model.args, list):
            raise TypeError("args must be a JSON list")
        if not isinstance(model.kwargs, dict):
            raise TypeError("kwargs must be a JSON object")
        if not isinstance(model.options, dict):
            raise TypeError("options must be a JSON object")
        if model.intervalSeconds < 1:
            raise ValueError("intervalSeconds must be greater than zero")

        return self.Entry(
            app=self.app,
            name=model.name,
            task=model.task,
            schedule=timedelta(seconds=model.intervalSeconds),
            args=tuple(model.args),
            kwargs=model.kwargs,
            options=model.options,
            last_run_at=model.lastRunAt,
            total_run_count=model.totalRunCount,
        )

    @staticmethod
    def _model_snapshot(model):
        return (
            model.pk,
            model.name,
            model.task,
            model.intervalSeconds,
            model.enabled,
            model.lastRunAt,
            model.totalRunCount,
            json.dumps(model.args, sort_keys=True, separators=(",", ":")),
            json.dumps(model.kwargs, sort_keys=True, separators=(",", ":")),
            json.dumps(model.options, sort_keys=True, separators=(",", ":")),
        )