final/backend/services/scheduler_service.py

"""
Enhanced Scheduler Service
Manages periodic and real-time data updates with persistence
"""
import asyncio
import logging
from typing import Dict, Any, List, Optional, Callable
from datetime import datetime, timedelta
from dataclasses import dataclass, asdict
import json
from collections import defaultdict
import httpx

logger = logging.getLogger(__name__)


@dataclass
class ScheduleTask:
    """Represents a scheduled task"""
    api_id: str
    name: str
    category: str
    interval: int  # seconds
    update_type: str  # realtime, periodic, scheduled
    enabled: bool
    last_update: Optional[datetime] = None
    next_update: Optional[datetime] = None
    last_status: Optional[str] = None  # success, failed, pending
    last_data: Optional[Dict[str, Any]] = None
    error_count: int = 0
    success_count: int = 0


class SchedulerService:
    """Advanced scheduler for managing API data updates"""

    def __init__(self, config_loader, db_manager=None):
        self.config_loader = config_loader
        self.db_manager = db_manager
        self.tasks: Dict[str, ScheduleTask] = {}
        self.running = False
        self.periodic_task = None
        self.realtime_tasks: Dict[str, asyncio.Task] = {}
        self.data_cache: Dict[str, Any] = {}
        self.callbacks: Dict[str, List[Callable]] = defaultdict(list)

        # Initialize tasks from config
        self._initialize_tasks()

    def _initialize_tasks(self):
        """Initialize schedule tasks from config loader"""
        apis = self.config_loader.get_all_apis()
        schedules = self.config_loader.schedules

        for api_id, api in apis.items():
            schedule = schedules.get(api_id, {})

            task = ScheduleTask(
                api_id=api_id,
                name=api.get('name', api_id),
                category=api.get('category', 'unknown'),
                interval=schedule.get('interval', 300),
                update_type=api.get('update_type', 'periodic'),
                enabled=schedule.get('enabled', True),
                next_update=datetime.now()
            )

            self.tasks[api_id] = task

        logger.info(f"Initialized {len(self.tasks)} schedule tasks")

    async def start(self):
        """Start the scheduler"""
        if self.running:
            logger.warning("Scheduler already running")
            return

        self.running = True
        logger.info("Starting scheduler...")

        # Start periodic update loop
        self.periodic_task = asyncio.create_task(self._periodic_update_loop())

        # Start real-time tasks
        await self._start_realtime_tasks()

        logger.info("Scheduler started successfully")

    async def stop(self):
        """Stop the scheduler"""
        if not self.running:
            return

        self.running = False
        logger.info("Stopping scheduler...")

        # Cancel periodic task
        if self.periodic_task:
            self.periodic_task.cancel()
            try:
                await self.periodic_task
            except asyncio.CancelledError:
                pass

        # Cancel real-time tasks
        for task in self.realtime_tasks.values():
            task.cancel()

        logger.info("Scheduler stopped")

    async def _periodic_update_loop(self):
        """Main loop for periodic updates"""
        while self.running:
            try:
                # Get tasks due for update
                due_tasks = self._get_due_tasks()

                if due_tasks:
                    logger.info(f"Processing {len(due_tasks)} due tasks")

                    # Process tasks concurrently
                    await asyncio.gather(
                        *[self._execute_task(task) for task in due_tasks],
                        return_exceptions=True
                    )

                # Sleep for a short interval
                await asyncio.sleep(5)  # Check every 5 seconds

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error in periodic update loop: {e}")
                await asyncio.sleep(10)

    def _get_due_tasks(self) -> List[ScheduleTask]:
        """Get tasks that are due for update"""
        now = datetime.now()
        due_tasks = []

        for task in self.tasks.values():
            if not task.enabled:
                continue

            if task.update_type == 'realtime':
                continue  # Real-time tasks handled separately

            if task.next_update is None or now >= task.next_update:
                due_tasks.append(task)

        return due_tasks

    async def _execute_task(self, task: ScheduleTask):
        """Execute a single scheduled task"""
        try:
            api = self.config_loader.apis.get(task.api_id)
            if not api:
                logger.error(f"API not found: {task.api_id}")
                return

            # Fetch data from API
            data = await self._fetch_api_data(api)

            # Update task status
            task.last_update = datetime.now()
            task.next_update = task.last_update + timedelta(seconds=task.interval)
            task.last_status = 'success'
            task.last_data = data
            task.success_count += 1
            task.error_count = 0  # Reset error count on success

            # Cache data
            self.data_cache[task.api_id] = {
                'data': data,
                'timestamp': datetime.now(),
                'task': task.name
            }

            # Save to database if available
            if self.db_manager:
                await self._save_to_database(task, data)

            # Trigger callbacks
            await self._trigger_callbacks(task.api_id, data)

            # Mark as updated in config loader
            self.config_loader.mark_updated(task.api_id)

            logger.info(f"✓ Updated {task.name} ({task.category})")

        except Exception as e:
            logger.error(f"✗ Failed to update {task.name}: {e}")
            task.last_status = 'failed'
            task.error_count += 1

            # Increase interval on repeated failures
            if task.error_count >= 3:
                task.interval = min(task.interval * 2, 3600)  # Max 1 hour
                logger.warning(f"Increased interval for {task.name} to {task.interval}s")

    async def _fetch_api_data(self, api: Dict[str, Any]) -> Dict[str, Any]:
        """Fetch data from an API"""
        base_url = api.get('base_url', '')
        auth = api.get('auth', {})

        # Build request URL
        url = base_url

        # Handle authentication
        headers = {}
        params = {}

        auth_type = auth.get('type', 'none')

        if auth_type == 'apiKey' or auth_type == 'apiKeyHeader':
            key = auth.get('key')
            header_name = auth.get('header_name', 'X-API-Key')
            if key:
                headers[header_name] = key

        elif auth_type == 'apiKeyQuery':
            key = auth.get('key')
            param_name = auth.get('param_name', 'apikey')
            if key:
                params[param_name] = key

        elif auth_type == 'apiKeyPath':
            key = auth.get('key')
            param_name = auth.get('param_name', 'API_KEY')
            if key:
                url = url.replace(f'{{{param_name}}}', key)

        # Make request
        timeout = httpx.Timeout(10.0)

        async with httpx.AsyncClient(timeout=timeout) as client:
            # Handle different endpoints
            endpoints = api.get('endpoints')

            if isinstance(endpoints, dict) and 'health' in endpoints:
                url = endpoints['health']
            elif isinstance(endpoints, str):
                url = endpoints

            # Add query params
            if params:
                url = f"{url}{'&' if '?' in url else '?'}" + '&'.join(f"{k}={v}" for k, v in params.items())

            response = await client.get(url, headers=headers)
            response.raise_for_status()

            return response.json()

    async def _save_to_database(self, task: ScheduleTask, data: Dict[str, Any]):
        """Save task data to database"""
        if not self.db_manager:
            return

        try:
            # Save using database manager
            await self.db_manager.save_collection_data(
                api_id=task.api_id,
                category=task.category,
                data=data,
                timestamp=datetime.now()
            )
        except Exception as e:
            logger.error(f"Error saving to database: {e}")

    async def _trigger_callbacks(self, api_id: str, data: Dict[str, Any]):
        """Trigger callbacks for API updates"""
        if api_id in self.callbacks:
            for callback in self.callbacks[api_id]:
                try:
                    if asyncio.iscoroutinefunction(callback):
                        await callback(api_id, data)
                    else:
                        callback(api_id, data)
                except Exception as e:
                    logger.error(f"Error in callback for {api_id}: {e}")

    async def _start_realtime_tasks(self):
        """Start WebSocket connections for real-time APIs"""
        realtime_apis = self.config_loader.get_realtime_apis()

        for api_id, api in realtime_apis.items():
            task = self.tasks.get(api_id)

            if task and task.enabled:
                # Create WebSocket task
                ws_task = asyncio.create_task(self._realtime_task(task, api))
                self.realtime_tasks[api_id] = ws_task

        logger.info(f"Started {len(self.realtime_tasks)} real-time tasks")

    async def _realtime_task(self, task: ScheduleTask, api: Dict[str, Any]):
        """Handle real-time WebSocket connection"""
        # This is a placeholder - implement WebSocket connection logic
        # based on the specific API requirements
        while self.running:
            try:
                # Connect to WebSocket
                # ws_url = api.get('base_url')
                # async with websockets.connect(ws_url) as ws:
                #     async for message in ws:
                #         data = json.loads(message)
                #         await self._handle_realtime_data(task, data)

                logger.info(f"Real-time task for {task.name} (placeholder)")
                await asyncio.sleep(60)  # Placeholder

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error in real-time task {task.name}: {e}")
                await asyncio.sleep(30)  # Retry after delay

    async def _handle_realtime_data(self, task: ScheduleTask, data: Dict[str, Any]):
        """Handle incoming real-time data"""
        task.last_update = datetime.now()
        task.last_status = 'success'
        task.last_data = data
        task.success_count += 1

        # Cache data
        self.data_cache[task.api_id] = {
            'data': data,
            'timestamp': datetime.now(),
            'task': task.name
        }

        # Save to database
        if self.db_manager:
            await self._save_to_database(task, data)

        # Trigger callbacks
        await self._trigger_callbacks(task.api_id, data)

    def register_callback(self, api_id: str, callback: Callable):
        """Register a callback for API updates"""
        self.callbacks[api_id].append(callback)

    def unregister_callback(self, api_id: str, callback: Callable):
        """Unregister a callback"""
        if api_id in self.callbacks:
            self.callbacks[api_id] = [cb for cb in self.callbacks[api_id] if cb != callback]

    def update_task_schedule(self, api_id: str, interval: int = None, enabled: bool = None):
        """Update schedule for a task"""
        if api_id in self.tasks:
            task = self.tasks[api_id]

            if interval is not None:
                task.interval = interval
                self.config_loader.update_schedule(api_id, interval=interval)

            if enabled is not None:
                task.enabled = enabled
                self.config_loader.update_schedule(api_id, enabled=enabled)

            logger.info(f"Updated schedule for {task.name}")

    def get_task_status(self, api_id: str) -> Optional[Dict[str, Any]]:
        """Get status of a specific task"""
        task = self.tasks.get(api_id)

        if not task:
            return None

        return {
            'api_id': task.api_id,
            'name': task.name,
            'category': task.category,
            'interval': task.interval,
            'update_type': task.update_type,
            'enabled': task.enabled,
            'last_update': task.last_update.isoformat() if task.last_update else None,
            'next_update': task.next_update.isoformat() if task.next_update else None,
            'last_status': task.last_status,
            'success_count': task.success_count,
            'error_count': task.error_count
        }

    def get_all_task_statuses(self) -> Dict[str, Any]:
        """Get status of all tasks"""
        return {
            api_id: self.get_task_status(api_id)
            for api_id in self.tasks.keys()
        }

    def get_cached_data(self, api_id: str) -> Optional[Dict[str, Any]]:
        """Get cached data for an API"""
        return self.data_cache.get(api_id)

    def get_all_cached_data(self) -> Dict[str, Any]:
        """Get all cached data"""
        return self.data_cache

    async def force_update(self, api_id: str) -> bool:
        """Force an immediate update for an API"""
        task = self.tasks.get(api_id)

        if not task:
            logger.error(f"Task not found: {api_id}")
            return False

        logger.info(f"Forcing update for {task.name}")
        await self._execute_task(task)

        return task.last_status == 'success'

    def export_schedules(self, filepath: str):
        """Export schedules to JSON"""
        schedules_data = {
            api_id: {
                'name': task.name,
                'category': task.category,
                'interval': task.interval,
                'update_type': task.update_type,
                'enabled': task.enabled,
                'last_update': task.last_update.isoformat() if task.last_update else None,
                'success_count': task.success_count,
                'error_count': task.error_count
            }
            for api_id, task in self.tasks.items()
        }

        with open(filepath, 'w') as f:
            json.dump(schedules_data, f, indent=2)

        logger.info(f"Exported schedules to {filepath}")

    def import_schedules(self, filepath: str):
        """Import schedules from JSON"""
        with open(filepath, 'r') as f:
            schedules_data = json.load(f)

        for api_id, schedule_data in schedules_data.items():
            if api_id in self.tasks:
                task = self.tasks[api_id]
                task.interval = schedule_data.get('interval', task.interval)
                task.enabled = schedule_data.get('enabled', task.enabled)

        logger.info(f"Imported schedules from {filepath}")