bot_manager.py

"""
MasterCryptoFarmBot — Bot instance manager.
Manages lifecycle of multiple farming bot instances: registration,
start/stop, health monitoring, and configuration per bot.
"""

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


class BotStatus(Enum):
    REGISTERED = auto()
    RUNNING = auto()
    PAUSED = auto()
    STOPPED = auto()
    ERROR = auto()


@dataclass
class BotInstance:
    bot_id: str
    name: str
    module: str
    status: BotStatus = BotStatus.REGISTERED
    api_type: str = "Pyrogram"
    started_at: float = 0.0
    stopped_at: float = 0.0
    cycles_completed: int = 0
    last_error: str = ""
    config: Dict[str, Any] = field(default_factory=dict)

    @property
    def uptime_seconds(self) -> float:
        if self.status == BotStatus.RUNNING and self.started_at > 0:
            return time.time() - self.started_at
        if self.stopped_at > self.started_at > 0:
            return self.stopped_at - self.started_at
        return 0.0

    def to_dict(self) -> Dict[str, Any]:
        return {
            "bot_id": self.bot_id,
            "name": self.name,
            "module": self.module,
            "status": self.status.name,
            "api_type": self.api_type,
            "uptime_seconds": round(self.uptime_seconds, 2),
            "cycles_completed": self.cycles_completed,
            "last_error": self.last_error,
        }


class BotManager:
    """
    Central manager for multiple farming bot instances.
    Handles registration, lifecycle control, health checks, and per-bot config.
    """

    def __init__(self):
        self._bots: Dict[str, BotInstance] = {}
        self._max_concurrent: int = 10
        self._on_status_change: Optional[Callable[[BotInstance, BotStatus], None]] = None

    @property
    def bot_count(self) -> int:
        return len(self._bots)

    @property
    def running_count(self) -> int:
        return sum(1 for b in self._bots.values() if b.status == BotStatus.RUNNING)

    def set_max_concurrent(self, limit: int) -> None:
        if limit < 1:
            raise ValueError("Concurrent limit must be at least 1")
        self._max_concurrent = limit

    def set_status_callback(
        self, callback: Callable[[BotInstance, BotStatus], None]
    ) -> None:
        self._on_status_change = callback

    def register(
        self,
        bot_id: str,
        name: str,
        module: str,
        api_type: str = "Pyrogram",
        config: Optional[Dict[str, Any]] = None,
    ) -> BotInstance:
        """Register a new farming bot instance."""
        if bot_id in self._bots:
            raise ValueError(f"Bot '{bot_id}' is already registered")
        instance = BotInstance(
            bot_id=bot_id,
            name=name,
            module=module,
            api_type=api_type,
            config=config or {},
        )
        self._bots[bot_id] = instance
        return instance

    def unregister(self, bot_id: str) -> None:
        """Remove a bot from the registry. Must be stopped first."""
        bot = self._get_bot(bot_id)
        if bot.status == BotStatus.RUNNING:
            raise RuntimeError(f"Cannot unregister running bot '{bot_id}'. Stop it first.")
        del self._bots[bot_id]

    def start(self, bot_id: str) -> BotInstance:
        """Start a registered bot instance."""
        bot = self._get_bot(bot_id)
        if bot.status == BotStatus.RUNNING:
            return bot
        if self.running_count >= self._max_concurrent:
            raise RuntimeError(
                f"Cannot start '{bot_id}': concurrent limit ({self._max_concurrent}) reached"
            )
        old_status = bot.status
        bot.status = BotStatus.RUNNING
        bot.started_at = time.time()
        bot.last_error = ""
        self._notify(bot, old_status)
        return bot

    def stop(self, bot_id: str) -> BotInstance:
        """Stop a running bot instance."""
        bot = self._get_bot(bot_id)
        if bot.status != BotStatus.RUNNING:
            return bot
        old_status = bot.status
        bot.status = BotStatus.STOPPED
        bot.stopped_at = time.time()
        self._notify(bot, old_status)
        return bot

    def pause(self, bot_id: str) -> BotInstance:
        """Pause a running bot instance."""
        bot = self._get_bot(bot_id)
        if bot.status != BotStatus.RUNNING:
            raise RuntimeError(f"Bot '{bot_id}' is not running")
        old_status = bot.status
        bot.status = BotStatus.PAUSED
        self._notify(bot, old_status)
        return bot

    def resume(self, bot_id: str) -> BotInstance:
        """Resume a paused bot instance."""
        bot = self._get_bot(bot_id)
        if bot.status != BotStatus.PAUSED:
            raise RuntimeError(f"Bot '{bot_id}' is not paused")
        old_status = bot.status
        bot.status = BotStatus.RUNNING
        self._notify(bot, old_status)
        return bot

    def report_error(self, bot_id: str, error_message: str) -> None:
        """Report an error for a bot instance."""
        bot = self._get_bot(bot_id)
        old_status = bot.status
        bot.status = BotStatus.ERROR
        bot.last_error = error_message
        bot.stopped_at = time.time()
        self._notify(bot, old_status)

    def increment_cycle(self, bot_id: str) -> int:
        """Increment the completed cycle counter for a bot."""
        bot = self._get_bot(bot_id)
        bot.cycles_completed += 1
        return bot.cycles_completed

    def start_all(self) -> List[BotInstance]:
        """Start all registered (non-running) bots up to the concurrent limit."""
        started: List[BotInstance] = []
        for bot_id, bot in self._bots.items():
            if bot.status != BotStatus.RUNNING and self.running_count < self._max_concurrent:
                started.append(self.start(bot_id))
        return started

    def stop_all(self) -> List[BotInstance]:
        """Stop all running bot instances."""
        stopped: List[BotInstance] = []
        for bot_id, bot in list(self._bots.items()):
            if bot.status == BotStatus.RUNNING:
                stopped.append(self.stop(bot_id))
        return stopped

    def get_bot(self, bot_id: str) -> Optional[BotInstance]:
        return self._bots.get(bot_id)

    def list_bots(self, status_filter: Optional[BotStatus] = None) -> List[BotInstance]:
        """List all bots, optionally filtered by status."""
        if status_filter is None:
            return list(self._bots.values())
        return [b for b in self._bots.values() if b.status == status_filter]

    def get_health_report(self) -> Dict[str, Any]:
        """Generate a health summary across all managed bots."""
        status_counts: Dict[str, int] = {}
        total_cycles = 0
        total_uptime = 0.0
        for bot in self._bots.values():
            key = bot.status.name
            status_counts[key] = status_counts.get(key, 0) + 1
            total_cycles += bot.cycles_completed
            total_uptime += bot.uptime_seconds
        return {
            "total_bots": self.bot_count,
            "status_distribution": status_counts,
            "total_cycles": total_cycles,
            "total_uptime_seconds": round(total_uptime, 2),
        }

    def _get_bot(self, bot_id: str) -> BotInstance:
        if bot_id not in self._bots:
            raise KeyError(f"Bot '{bot_id}' not found")
        return self._bots[bot_id]

    def _notify(self, bot: BotInstance, old_status: BotStatus) -> None:
        if self._on_status_change and bot.status != old_status:
            self._on_status_change(bot, old_status)