import asyncio
import random
import time
from typing import Callable

from loguru import logger
from pydantic import BaseModel

from common.scheduling.job_handler import Job, LOW_PRIORITY_JOB_QUEUE_NAME
from common.scheduling.queue_utils import (
    get_key,
    set_key,
    attempt_lock_capture_with_token,
    release_lock_with_token,
)


class PeriodicJob(BaseModel):
    """A Job with periodic scheduling parameters."""

    job: Job
    interval_seconds: float
    run_on_startup: bool = True  # If False, skip first run on system startup
    jitter_seconds: float = (
        0  # Random delay added to interval to prevent thundering herd
    )
    enabled: bool = True  # If False, job is registered but won't be scheduled

    def identify(self) -> str:
        return self.job.identify()

    def get_effective_interval(self) -> float:
        """Get interval with random jitter applied."""
        if self.jitter_seconds <= 0:
            return self.interval_seconds
        return self.interval_seconds + random.uniform(0, self.jitter_seconds)


class GlobalBeatService:
    """
    Beat service for spawning periodic jobs.

    Configure in session_manager.py:
        from common.scheduling.beat import GlobalBeatService
        BeatService = GlobalBeatService

    Register jobs with decorator:
        class MaintenanceTasks:
            @staticmethod
            @BeatService.run_every(interval_seconds=3600)
            async def cleanup():
                ...

    Run alongside consumers:
        await asyncio.gather(
            consume_all(multiplicity=10),
            BeatService.run(),
        )

    Constraint: Decorated functions must be class methods compatible with
    Job.new(). The executor provides @SessionManager.with_sync_session,
    so decorated methods don't need their own session decorator.
    """

    _jobs: list[PeriodicJob] = []
    _last_run_key_prefix: str = "beat:last_run:"

    @classmethod
    def run_every(
        cls,
        interval_seconds: float,
        queue_name: str = LOW_PRIORITY_JOB_QUEUE_NAME,
        run_on_startup: bool = True,
        jitter_seconds: float = 0,
        enabled: bool = True,
    ) -> Callable:
        """
        Decorator to register a classmethod/staticmethod as a periodic job.

        Args:
            interval_seconds: How often the job should run.
            queue_name: Which queue to enqueue to (default: low_priority).
            run_on_startup: If True (default), job runs immediately on first startup.
                If False, waits for the interval to elapse before first run.
            jitter_seconds: Random delay (0 to jitter_seconds) added to interval
                to prevent thundering herd across replicas.
            enabled: If False, job is registered but won't be scheduled.
        """

        def decorator(fn: Callable) -> Callable:
            job = Job.new(fn, queue_name=queue_name, ensure_sql_model=False)
            cls._jobs.append(
                PeriodicJob(
                    job=job,
                    interval_seconds=interval_seconds,
                    run_on_startup=run_on_startup,
                    jitter_seconds=jitter_seconds,
                    enabled=enabled,
                )
            )
            logger.debug(
                f"Beat registered: {job.identify()} (every {interval_seconds}s, "
                f"jitter={jitter_seconds}s, run_on_startup={run_on_startup}, "
                f"enabled={enabled})"
            )
            return fn

        return decorator

    @classmethod
    def register(
        cls,
        job_fn: str | Callable,
        interval_seconds: float,
        queue_name: str = LOW_PRIORITY_JOB_QUEUE_NAME,
        run_on_startup: bool = True,
        jitter_seconds: float = 0,
        enabled: bool = True,
        **job_kwargs,
    ) -> PeriodicJob:
        """
        Programmatically register a periodic job.

        Args:
            job_fn: Either a callable (classmethod/staticmethod) or a string like
                "module.Class.method" that Job.new() can resolve.
            interval_seconds: How often the job should run.
            queue_name: Which queue to enqueue to (default: low_priority).
            run_on_startup: If True (default), job runs immediately on first startup.
                If False, waits for the interval to elapse before first run.
            jitter_seconds: Random delay (0 to jitter_seconds) added to interval
                to prevent thundering herd across replicas.
            enabled: If False, job is registered but won't be scheduled.
            **job_kwargs: Additional keyword arguments to pass to the job when executed.

        Returns:
            The registered PeriodicJob.
        """
        job = Job.new(
            job_fn, queue_name=queue_name, ensure_sql_model=False, **job_kwargs
        )
        periodic_job = PeriodicJob(
            job=job,
            interval_seconds=interval_seconds,
            run_on_startup=run_on_startup,
            jitter_seconds=jitter_seconds,
            enabled=enabled,
        )
        cls._jobs.append(periodic_job)
        logger.debug(
            f"Beat registered: {job.identify()} (every {interval_seconds}s, "
            f"jitter={jitter_seconds}s, run_on_startup={run_on_startup}, "
            f"enabled={enabled})"
        )
        return periodic_job

    @classmethod
    def list_jobs(cls) -> list[PeriodicJob]:
        """List all registered periodic jobs."""
        return cls._jobs.copy()

    @classmethod
    async def run(cls, tick_interval_seconds: float = 60):
        """
        Run the beat service loop.

        Args:
            tick_interval_seconds: How often to check if jobs are due.
        """
        enabled_jobs = [job for job in cls._jobs if job.enabled]
        disabled_jobs = [job for job in cls._jobs if not job.enabled]
        logger.info(
            f"Beat service starting with {len(enabled_jobs)} enabled periodic jobs "
            f"({len(disabled_jobs)} disabled)"
        )
        for job in cls._jobs:
            jitter_info = (
                f", jitter={job.jitter_seconds}s" if job.jitter_seconds > 0 else ""
            )
            status = "" if job.enabled else " [DISABLED]"
            logger.info(
                f"  - {job.identify()} (every {job.interval_seconds}s{jitter_info}){status}"
            )

        while True:
            now = time.time()

            for periodic_job in cls._jobs:
                try:
                    await cls._maybe_spawn(periodic_job, now)
                except Exception:
                    logger.exception(
                        f"Beat failed to check job: {periodic_job.identify()}"
                    )

            await asyncio.sleep(tick_interval_seconds)

    @classmethod
    async def _maybe_spawn(cls, periodic_job: PeriodicJob, now: float):
        """Spawn job if interval has elapsed. Uses atomic lock to prevent duplicate spawns."""
        if not periodic_job.enabled:
            return  # Job is disabled, skip

        job_id = periodic_job.identify()
        key = f"{cls._last_run_key_prefix}{job_id}"

        # Check if enough time has passed
        last_run = await get_key(key)

        # Handle run_on_startup flag: if False and no last_run exists,
        # initialize to now so the job waits for the interval before first run
        if last_run is None:
            if not periodic_job.run_on_startup:
                await set_key(key, str(now))
                logger.debug(
                    f"Beat initialized last_run for {job_id} (run_on_startup=False)"
                )
                return  # Skip this tick, will run after interval elapses
            last_run_time = 0
        else:
            last_run_time = float(last_run)

        if now - last_run_time < periodic_job.get_effective_interval():
            return  # Not due yet

        # Atomic lock to prevent duplicate spawns across replicas
        lock_key = f"{key}:lock"
        acquired, lock_token = await attempt_lock_capture_with_token(
            lock_key, expire_in_seconds=30
        )

        if not acquired:
            return  # Another replica is spawning this job

        try:
            # Enqueue first, then update last_run on success
            # This ensures we retry on next interval if enqueue fails
            await periodic_job.job.enqueue()
            await set_key(key, str(now))
            logger.info(f"Beat spawned: {job_id}")
        finally:
            await release_lock_with_token(lock_key, lock_token)
