scheduler.py

"""
MasterCryptoFarmBot — Task scheduler and rotation engine.
Manages scheduling of farming tasks across multiple bot instances
with configurable intervals, priorities, and rotation strategies.
"""

import heapq
import time
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Any, Callable, Dict, List, Optional


class TaskPriority(Enum):
    LOW = 3
    NORMAL = 2
    HIGH = 1
    CRITICAL = 0


class TaskState(Enum):
    QUEUED = auto()
    RUNNING = auto()
    COMPLETED = auto()
    FAILED = auto()
    SKIPPED = auto()


class RotationStrategy(Enum):
    ROUND_ROBIN = auto()
    RANDOM = auto()
    PRIORITY = auto()
    WEIGHTED = auto()


@dataclass(order=True)
class ScheduledTask:
    priority_value: int = field(compare=True)
    scheduled_at: float = field(compare=True)
    task_id: str = field(compare=False)
    bot_id: str = field(compare=False)
    action_name: str = field(compare=False)
    state: TaskState = field(default=TaskState.QUEUED, compare=False)
    interval_seconds: float = field(default=0.0, compare=False)
    recurring: bool = field(default=False, compare=False)
    max_retries: int = field(default=3, compare=False)
    retries: int = field(default=0, compare=False)
    last_run: float = field(default=0.0, compare=False)
    last_result: str = field(default="", compare=False)
    metadata: Dict[str, Any] = field(default_factory=dict, compare=False)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "task_id": self.task_id,
            "bot_id": self.bot_id,
            "action_name": self.action_name,
            "state": self.state.name,
            "priority": self.priority_value,
            "scheduled_at": self.scheduled_at,
            "interval_seconds": self.interval_seconds,
            "recurring": self.recurring,
            "retries": self.retries,
            "last_run": self.last_run,
            "last_result": self.last_result,
        }


class TaskScheduler:
    """
    Priority-queue-based task scheduler for farming operations.
    Supports recurring tasks, rotation strategies, and per-bot task assignment.
    """

    def __init__(self, strategy: RotationStrategy = RotationStrategy.PRIORITY):
        self._queue: List[ScheduledTask] = []
        self._completed: List[ScheduledTask] = []
        self._strategy = strategy
        self._task_counter = 0
        self._paused = False
        self._bot_round_robin_index: Dict[str, int] = {}

    @property
    def pending_count(self) -> int:
        return len(self._queue)

    @property
    def completed_count(self) -> int:
        return len(self._completed)

    @property
    def is_paused(self) -> bool:
        return self._paused

    def schedule(
        self,
        bot_id: str,
        action_name: str,
        priority: TaskPriority = TaskPriority.NORMAL,
        delay_seconds: float = 0.0,
        interval_seconds: float = 0.0,
        recurring: bool = False,
        max_retries: int = 3,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> ScheduledTask:
        """Schedule a new task for a specific bot."""
        self._task_counter += 1
        task = ScheduledTask(
            priority_value=priority.value,
            scheduled_at=time.time() + delay_seconds,
            task_id=f"task-{self._task_counter:06d}",
            bot_id=bot_id,
            action_name=action_name,
            interval_seconds=interval_seconds,
            recurring=recurring,
            max_retries=max_retries,
            metadata=metadata or {},
        )
        heapq.heappush(self._queue, task)
        return task

    def schedule_batch(
        self,
        bot_ids: List[str],
        action_name: str,
        priority: TaskPriority = TaskPriority.NORMAL,
        stagger_seconds: float = 1.0,
    ) -> List[ScheduledTask]:
        """Schedule the same action across multiple bots with staggered start."""
        tasks: List[ScheduledTask] = []
        for i, bid in enumerate(bot_ids):
            task = self.schedule(
                bot_id=bid,
                action_name=action_name,
                priority=priority,
                delay_seconds=i * stagger_seconds,
            )
            tasks.append(task)
        return tasks

    def next_task(self) -> Optional[ScheduledTask]:
        """Dequeue the highest-priority task that is ready to run."""
        if self._paused or not self._queue:
            return None

        now = time.time()
        ready: List[ScheduledTask] = []
        remaining: List[ScheduledTask] = []

        while self._queue:
            task = heapq.heappop(self._queue)
            if task.scheduled_at <= now and task.state == TaskState.QUEUED:
                ready.append(task)
            else:
                remaining.append(task)

        for t in remaining:
            heapq.heappush(self._queue, t)

        if not ready:
            return None

        selected = ready[0]
        selected.state = TaskState.RUNNING
        selected.last_run = now

        for t in ready[1:]:
            heapq.heappush(self._queue, t)

        return selected

    def complete_task(self, task: ScheduledTask, result: str = "success") -> None:
        """Mark a task as completed and optionally reschedule if recurring."""
        task.state = TaskState.COMPLETED
        task.last_result = result
        self._completed.append(task)

        if task.recurring and task.interval_seconds > 0:
            new_task = ScheduledTask(
                priority_value=task.priority_value,
                scheduled_at=time.time() + task.interval_seconds,
                task_id=task.task_id + "-r",
                bot_id=task.bot_id,
                action_name=task.action_name,
                interval_seconds=task.interval_seconds,
                recurring=True,
                max_retries=task.max_retries,
                metadata=task.metadata,
            )
            heapq.heappush(self._queue, new_task)

    def fail_task(self, task: ScheduledTask, error: str = "") -> None:
        """Mark a task as failed, retrying if under the retry limit."""
        task.retries += 1
        task.last_result = error

        if task.retries < task.max_retries:
            task.state = TaskState.QUEUED
            task.scheduled_at = time.time() + (2 ** task.retries)
            heapq.heappush(self._queue, task)
        else:
            task.state = TaskState.FAILED
            self._completed.append(task)

    def skip_task(self, task: ScheduledTask, reason: str = "") -> None:
        task.state = TaskState.SKIPPED
        task.last_result = reason
        self._completed.append(task)

    def pause(self) -> None:
        self._paused = True

    def resume(self) -> None:
        self._paused = False

    def clear_queue(self) -> int:
        count = len(self._queue)
        self._queue.clear()
        return count

    def get_tasks_for_bot(self, bot_id: str) -> List[ScheduledTask]:
        return [t for t in self._queue if t.bot_id == bot_id]

    def get_completed_for_bot(self, bot_id: str) -> List[ScheduledTask]:
        return [t for t in self._completed if t.bot_id == bot_id]

    def get_queue_snapshot(self) -> List[Dict[str, Any]]:
        return sorted([t.to_dict() for t in self._queue], key=lambda x: x["scheduled_at"])

    def get_stats(self) -> Dict[str, Any]:
        failed = sum(1 for t in self._completed if t.state == TaskState.FAILED)
        completed = sum(1 for t in self._completed if t.state == TaskState.COMPLETED)
        skipped = sum(1 for t in self._completed if t.state == TaskState.SKIPPED)
        return {
            "pending": self.pending_count,
            "completed": completed,
            "failed": failed,
            "skipped": skipped,
            "total_processed": len(self._completed),
            "strategy": self._strategy.name,
            "paused": self._paused,
        }