monitoring/source_pool_manager.py

"""
Intelligent Source Pool Manager
Manages source pools, rotation, and automatic failover
"""

import json
from datetime import datetime, timedelta
from typing import Optional, List, Dict, Any
from threading import Lock
from sqlalchemy.orm import Session

from database.models import (
    SourcePool, PoolMember, RotationHistory, RotationState,
    Provider, RateLimitUsage
)
from monitoring.rate_limiter import rate_limiter
from utils.logger import setup_logger

logger = setup_logger("source_pool_manager")


class SourcePoolManager:
    """
    Manages source pools and intelligent rotation
    """

    def __init__(self, db_session: Session):
        """
        Initialize source pool manager

        Args:
            db_session: Database session
        """
        self.db = db_session
        self.lock = Lock()
        logger.info("Source Pool Manager initialized")

    def create_pool(
        self,
        name: str,
        category: str,
        description: Optional[str] = None,
        rotation_strategy: str = "round_robin"
    ) -> SourcePool:
        """
        Create a new source pool

        Args:
            name: Pool name
            category: Pool category
            description: Pool description
            rotation_strategy: Rotation strategy (round_robin, least_used, priority)

        Returns:
            Created SourcePool
        """
        with self.lock:
            pool = SourcePool(
                name=name,
                category=category,
                description=description,
                rotation_strategy=rotation_strategy,
                enabled=True
            )
            self.db.add(pool)
            self.db.commit()
            self.db.refresh(pool)

            # Create rotation state
            state = RotationState(
                pool_id=pool.id,
                current_provider_id=None,
                rotation_count=0
            )
            self.db.add(state)
            self.db.commit()

            logger.info(f"Created source pool: {name} (strategy: {rotation_strategy})")
            return pool

    def add_to_pool(
        self,
        pool_id: int,
        provider_id: int,
        priority: int = 1,
        weight: int = 1
    ) -> PoolMember:
        """
        Add a provider to a pool

        Args:
            pool_id: Pool ID
            provider_id: Provider ID
            priority: Provider priority (higher = better)
            weight: Provider weight for weighted rotation

        Returns:
            Created PoolMember
        """
        with self.lock:
            member = PoolMember(
                pool_id=pool_id,
                provider_id=provider_id,
                priority=priority,
                weight=weight,
                enabled=True,
                use_count=0,
                success_count=0,
                failure_count=0
            )
            self.db.add(member)
            self.db.commit()
            self.db.refresh(member)

            logger.info(f"Added provider {provider_id} to pool {pool_id}")
            return member

    def get_next_provider(
        self,
        pool_id: int,
        exclude_rate_limited: bool = True
    ) -> Optional[Provider]:
        """
        Get next provider from pool based on rotation strategy

        Args:
            pool_id: Pool ID
            exclude_rate_limited: Exclude rate-limited providers

        Returns:
            Next Provider or None if none available
        """
        with self.lock:
            # Get pool and its members
            pool = self.db.query(SourcePool).filter_by(id=pool_id).first()
            if not pool or not pool.enabled:
                logger.warning(f"Pool {pool_id} not found or disabled")
                return None

            # Get enabled members with their providers
            members = (
                self.db.query(PoolMember)
                .filter_by(pool_id=pool_id, enabled=True)
                .join(Provider)
                .filter(Provider.id == PoolMember.provider_id)
                .all()
            )

            if not members:
                logger.warning(f"No enabled members in pool {pool_id}")
                return None

            # Filter out rate-limited providers
            if exclude_rate_limited:
                available_members = []
                for member in members:
                    provider = self.db.query(Provider).get(member.provider_id)
                    can_use, _ = rate_limiter.can_make_request(provider.name)
                    if can_use:
                        available_members.append(member)

                if not available_members:
                    logger.warning(f"All providers in pool {pool_id} are rate-limited")
                    # Return highest priority member anyway
                    available_members = members
            else:
                available_members = members

            # Select provider based on strategy
            selected_member = self._select_by_strategy(
                pool.rotation_strategy,
                available_members
            )

            if not selected_member:
                return None

            # Get rotation state
            state = self.db.query(RotationState).filter_by(pool_id=pool_id).first()
            if not state:
                state = RotationState(pool_id=pool_id)
                self.db.add(state)

            # Record rotation if provider changed
            old_provider_id = state.current_provider_id
            if old_provider_id != selected_member.provider_id:
                self._record_rotation(
                    pool_id=pool_id,
                    from_provider_id=old_provider_id,
                    to_provider_id=selected_member.provider_id,
                    reason="rotation"
                )

            # Update state
            state.current_provider_id = selected_member.provider_id
            state.last_rotation = datetime.utcnow()
            state.rotation_count += 1

            # Update member stats
            selected_member.last_used = datetime.utcnow()
            selected_member.use_count += 1

            self.db.commit()

            provider = self.db.query(Provider).get(selected_member.provider_id)
            logger.info(
                f"Selected provider {provider.name} from pool {pool.name} "
                f"(strategy: {pool.rotation_strategy})"
            )
            return provider

    def _select_by_strategy(
        self,
        strategy: str,
        members: List[PoolMember]
    ) -> Optional[PoolMember]:
        """
        Select a pool member based on rotation strategy

        Args:
            strategy: Rotation strategy
            members: Available pool members

        Returns:
            Selected PoolMember
        """
        if not members:
            return None

        if strategy == "priority":
            # Select highest priority member
            return max(members, key=lambda m: m.priority)

        elif strategy == "least_used":
            # Select least used member
            return min(members, key=lambda m: m.use_count)

        elif strategy == "weighted":
            # Weighted random selection (simple implementation)
            # In production, use proper weighted random
            return max(members, key=lambda m: m.weight * (1.0 / (m.use_count + 1)))

        else:  # round_robin (default)
            # Select least recently used
            never_used = [m for m in members if m.last_used is None]
            if never_used:
                return never_used[0]
            return min(members, key=lambda m: m.last_used)

    def _record_rotation(
        self,
        pool_id: int,
        from_provider_id: Optional[int],
        to_provider_id: int,
        reason: str,
        notes: Optional[str] = None
    ):
        """
        Record a rotation event

        Args:
            pool_id: Pool ID
            from_provider_id: Previous provider ID
            to_provider_id: New provider ID
            reason: Rotation reason
            notes: Additional notes
        """
        rotation = RotationHistory(
            pool_id=pool_id,
            from_provider_id=from_provider_id,
            to_provider_id=to_provider_id,
            rotation_reason=reason,
            success=True,
            notes=notes
        )
        self.db.add(rotation)
        self.db.commit()

    def failover(
        self,
        pool_id: int,
        failed_provider_id: int,
        reason: str = "failure"
    ) -> Optional[Provider]:
        """
        Perform failover from a failed provider

        Args:
            pool_id: Pool ID
            failed_provider_id: Failed provider ID
            reason: Failure reason

        Returns:
            Next available provider
        """
        with self.lock:
            logger.warning(
                f"Failover triggered for provider {failed_provider_id} "
                f"in pool {pool_id}. Reason: {reason}"
            )

            # Update failure count for the failed provider
            member = (
                self.db.query(PoolMember)
                .filter_by(pool_id=pool_id, provider_id=failed_provider_id)
                .first()
            )
            if member:
                member.failure_count += 1
                self.db.commit()

            # Get next provider (excluding the failed one)
            pool = self.db.query(SourcePool).filter_by(id=pool_id).first()
            if not pool:
                return None

            members = (
                self.db.query(PoolMember)
                .filter_by(pool_id=pool_id, enabled=True)
                .filter(PoolMember.provider_id != failed_provider_id)
                .all()
            )

            if not members:
                logger.error(f"No alternative providers available in pool {pool_id}")
                return None

            # Select next provider
            selected_member = self._select_by_strategy(
                pool.rotation_strategy,
                members
            )

            if not selected_member:
                return None

            # Record failover
            self._record_rotation(
                pool_id=pool_id,
                from_provider_id=failed_provider_id,
                to_provider_id=selected_member.provider_id,
                reason=reason,
                notes=f"Automatic failover from provider {failed_provider_id}"
            )

            # Update rotation state
            state = self.db.query(RotationState).filter_by(pool_id=pool_id).first()
            if state:
                state.current_provider_id = selected_member.provider_id
                state.last_rotation = datetime.utcnow()
                state.rotation_count += 1

            # Update member stats
            selected_member.last_used = datetime.utcnow()
            selected_member.use_count += 1

            self.db.commit()

            provider = self.db.query(Provider).get(selected_member.provider_id)
            logger.info(f"Failover successful: switched to provider {provider.name}")
            return provider

    def record_success(self, pool_id: int, provider_id: int):
        """
        Record successful use of a provider

        Args:
            pool_id: Pool ID
            provider_id: Provider ID
        """
        with self.lock:
            member = (
                self.db.query(PoolMember)
                .filter_by(pool_id=pool_id, provider_id=provider_id)
                .first()
            )
            if member:
                member.success_count += 1
                self.db.commit()

    def record_failure(self, pool_id: int, provider_id: int):
        """
        Record failed use of a provider

        Args:
            pool_id: Pool ID
            provider_id: Provider ID
        """
        with self.lock:
            member = (
                self.db.query(PoolMember)
                .filter_by(pool_id=pool_id, provider_id=provider_id)
                .first()
            )
            if member:
                member.failure_count += 1
                self.db.commit()

    def get_pool_status(self, pool_id: int) -> Optional[Dict[str, Any]]:
        """
        Get comprehensive pool status

        Args:
            pool_id: Pool ID

        Returns:
            Pool status dictionary
        """
        with self.lock:
            pool = self.db.query(SourcePool).filter_by(id=pool_id).first()
            if not pool:
                return None

            # Get rotation state
            state = self.db.query(RotationState).filter_by(pool_id=pool_id).first()

            # Get current provider
            current_provider = None
            if state and state.current_provider_id:
                provider = self.db.query(Provider).get(state.current_provider_id)
                if provider:
                    current_provider = {
                        "id": provider.id,
                        "name": provider.name,
                        "status": "active"
                    }

            # Get all members with stats
            members = []
            pool_members = self.db.query(PoolMember).filter_by(pool_id=pool_id).all()

            for member in pool_members:
                provider = self.db.query(Provider).get(member.provider_id)
                if not provider:
                    continue

                # Check rate limit status
                rate_status = rate_limiter.get_status(provider.name)
                rate_limit_info = None
                if rate_status:
                    rate_limit_info = {
                        "usage": rate_status['current_usage'],
                        "limit": rate_status['limit_value'],
                        "percentage": rate_status['percentage'],
                        "status": rate_status['status']
                    }

                success_rate = 0
                if member.use_count > 0:
                    success_rate = (member.success_count / member.use_count) * 100

                members.append({
                    "provider_id": provider.id,
                    "provider_name": provider.name,
                    "priority": member.priority,
                    "weight": member.weight,
                    "enabled": member.enabled,
                    "use_count": member.use_count,
                    "success_count": member.success_count,
                    "failure_count": member.failure_count,
                    "success_rate": round(success_rate, 2),
                    "last_used": member.last_used.isoformat() if member.last_used else None,
                    "rate_limit": rate_limit_info
                })

            # Get recent rotations
            recent_rotations = (
                self.db.query(RotationHistory)
                .filter_by(pool_id=pool_id)
                .order_by(RotationHistory.timestamp.desc())
                .limit(10)
                .all()
            )

            rotation_list = []
            for rotation in recent_rotations:
                from_provider = None
                if rotation.from_provider_id:
                    from_prov = self.db.query(Provider).get(rotation.from_provider_id)
                    from_provider = from_prov.name if from_prov else None

                to_prov = self.db.query(Provider).get(rotation.to_provider_id)
                to_provider = to_prov.name if to_prov else None

                rotation_list.append({
                    "timestamp": rotation.timestamp.isoformat(),
                    "from_provider": from_provider,
                    "to_provider": to_provider,
                    "reason": rotation.rotation_reason,
                    "success": rotation.success
                })

            return {
                "pool_id": pool.id,
                "pool_name": pool.name,
                "category": pool.category,
                "description": pool.description,
                "rotation_strategy": pool.rotation_strategy,
                "enabled": pool.enabled,
                "current_provider": current_provider,
                "total_rotations": state.rotation_count if state else 0,
                "last_rotation": state.last_rotation.isoformat() if state and state.last_rotation else None,
                "members": members,
                "recent_rotations": rotation_list
            }

    def get_all_pools_status(self) -> List[Dict[str, Any]]:
        """
        Get status of all pools

        Returns:
            List of pool status dictionaries
        """
        pools = self.db.query(SourcePool).all()
        return [
            self.get_pool_status(pool.id)
            for pool in pools
            if self.get_pool_status(pool.id)
        ]