"""Auto-scaling worker pool management for OCLAWMA.

This module provides automatic scaling of worker pools based on queue depth,
processing rate, and configurable thresholds. Supports Kubernetes HPA and
cloud provider auto-scaling groups.
"""

from __future__ import annotations

import logging
import threading
import time
from contextlib import suppress
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Any, Protocol

logger = logging.getLogger(__name__)


class ScalingDecision(Enum):
    """Decision from the scaling recommendation."""

    SCALE_UP = auto()
    SCALE_DOWN = auto()
    MAINTAIN = auto()


@dataclass
class ScalingThresholds:
    """Configuration thresholds for auto-scaling.

    Attributes:
        scale_up_threshold: Queue depth to trigger scale-up
        scale_down_threshold: Queue depth to trigger scale-down
        max_workers: Maximum number of workers allowed
        min_workers: Minimum number of workers allowed
        scale_up_cooldown: Seconds to wait between scale-up operations
        scale_down_cooldown: Seconds to wait between scale-down operations
        processing_rate_window: Seconds of history to consider for rate calculation
    """

    scale_up_threshold: int = 100
    scale_down_threshold: int = 10
    max_workers: int = 10
    min_workers: int = 1
    scale_up_cooldown: float = 60.0
    scale_down_cooldown: float = 120.0
    processing_rate_window: float = 300.0  # 5 minutes

    def __post_init__(self) -> None:
        """Validate threshold configuration."""
        if self.min_workers < 0:
            raise ValueError("min_workers must be non-negative")
        if self.max_workers < self.min_workers:
            raise ValueError("max_workers must be >= min_workers")
        if self.scale_up_threshold <= self.scale_down_threshold:
            raise ValueError("scale_up_threshold must be > scale_down_threshold")
        if self.scale_up_cooldown < 0 or self.scale_down_cooldown < 0:
            raise ValueError("cooldown periods must be non-negative")


@dataclass
class ScalingMetrics:
    """Current metrics for scaling decisions.

    Attributes:
        queue_depth: Current number of jobs in queue
        processing_rate: Jobs processed per second (moving average)
        worker_count: Current number of active workers
        timestamp: Time when metrics were recorded
    """

    queue_depth: int = 0
    processing_rate: float = 0.0
    worker_count: int = 0
    timestamp: float = field(default_factory=time.time)


@dataclass
class ScalingRecommendation:
    """Recommendation from the auto-scaler.

    Attributes:
        decision: The scaling decision (up, down, or maintain)
        target_workers: Recommended number of workers
        current_workers: Current number of workers
        reason: Human-readable explanation for the decision
        metrics: Metrics that led to this recommendation
    """

    decision: ScalingDecision
    target_workers: int
    current_workers: int
    reason: str
    metrics: ScalingMetrics


class WorkerPoolBackend(Protocol):
    """Protocol for worker pool backend implementations.

    Implementations can target Kubernetes, cloud providers, or local processes.
    """

    def get_worker_count(self) -> int:
        """Get current number of workers."""
        ...

    def scale_up(self, count: int) -> bool:
        """Scale up by adding workers. Returns success status."""
        ...

    def scale_down(self, count: int) -> bool:
        """Scale down by removing workers. Returns success status."""
        ...

    def set_worker_count(self, count: int) -> bool:
        """Set exact worker count. Returns success status."""
        ...


class LocalWorkerBackend:
    """Local process-based worker backend for testing/development."""

    def __init__(self, initial_workers: int = 1) -> None:
        """Initialize local backend.

        Args:
            initial_workers: Starting number of workers
        """
        self._worker_count = initial_workers
        self._lock = threading.RLock()

    def get_worker_count(self) -> int:
        """Get current number of workers."""
        with self._lock:
            return self._worker_count

    def scale_up(self, count: int) -> bool:
        """Scale up by adding workers."""
        with self._lock:
            self._worker_count += count
            logger.info(f"Local backend: scaled up by {count} to {self._worker_count}")
            return True

    def scale_down(self, count: int) -> bool:
        """Scale down by removing workers."""
        with self._lock:
            new_count = max(0, self._worker_count - count)
            actual_removed = self._worker_count - new_count
            self._worker_count = new_count
            logger.info(f"Local backend: scaled down by {actual_removed} to {self._worker_count}")
            return True

    def set_worker_count(self, count: int) -> bool:
        """Set exact worker count."""
        with self._lock:
            self._worker_count = max(0, count)
            logger.info(f"Local backend: set worker count to {self._worker_count}")
            return True


class KubernetesHpaBackend:
    """Kubernetes HPA-based scaling backend.

    Requires kubernetes client library to be installed.
    """

    def __init__(
        self,
        deployment_name: str,
        namespace: str = "default",
        hpa_name: str | None = None,
    ) -> None:
        """Initialize Kubernetes backend.

        Args:
            deployment_name: Name of the Deployment to scale
            namespace: Kubernetes namespace
            hpa_name: Name of HPA resource (defaults to deployment_name)
        """
        self.deployment_name = deployment_name
        self.namespace = namespace
        self.hpa_name = hpa_name or deployment_name
        self._k8s_available = False
        self._client: Any = None

        try:
            from kubernetes import client, config

            config.load_incluster_config() if self._in_cluster() else config.load_kube_config()
            self._client = client
            self._k8s_available = True
        except ImportError:
            logger.warning("kubernetes client not installed, K8s backend disabled")
        except Exception as e:
            logger.warning(f"Failed to initialize Kubernetes client: {e}")

    def _in_cluster(self) -> bool:
        """Check if running inside Kubernetes cluster."""
        import os

        return os.path.exists("/var/run/secrets/kubernetes.io/serviceaccount/token")

    def get_worker_count(self) -> int:
        """Get current number of replicas from Deployment."""
        if not self._k8s_available:
            logger.warning("Kubernetes not available, returning 0")
            return 0

        try:
            api = self._client.AppsV1Api()
            deployment = api.read_namespaced_deployment(
                name=self.deployment_name,
                namespace=self.namespace,
            )
            return deployment.spec.replicas or 0
        except Exception as e:
            logger.error(f"Failed to get worker count: {e}")
            return 0

    def scale_up(self, count: int) -> bool:
        """Scale up by patching the HPA min replicas."""
        if not self._k8s_available:
            logger.error("Kubernetes not available")
            return False

        try:
            current = self.get_worker_count()
            return self.set_worker_count(current + count)
        except Exception as e:
            logger.error(f"Failed to scale up: {e}")
            return False

    def scale_down(self, count: int) -> bool:
        """Scale down by patching the HPA min replicas."""
        if not self._k8s_available:
            logger.error("Kubernetes not available")
            return False

        try:
            current = self.get_worker_count()
            return self.set_worker_count(max(0, current - count))
        except Exception as e:
            logger.error(f"Failed to scale down: {e}")
            return False

    def set_worker_count(self, count: int) -> bool:
        """Set exact replica count on the Deployment."""
        if not self._k8s_available:
            logger.error("Kubernetes not available")
            return False

        try:
            api = self._client.AppsV1Api()
            patch = {"spec": {"replicas": count}}
            api.patch_namespaced_deployment_scale(
                name=self.deployment_name,
                namespace=self.namespace,
                body=patch,
            )
            logger.info(f"Kubernetes: set {self.deployment_name} replicas to {count}")
            return True
        except Exception as e:
            logger.error(f"Failed to set worker count: {e}")
            return False


class AwsAsgBackend:
    """AWS Auto Scaling Group backend.

    Requires boto3 to be installed and AWS credentials configured.
    """

    def __init__(
        self,
        asg_name: str,
        region: str = "us-east-1",
    ) -> None:
        """Initialize AWS ASG backend.

        Args:
            asg_name: Name of the Auto Scaling Group
            region: AWS region
        """
        self.asg_name = asg_name
        self.region = region
        self._aws_available = False
        self._client: Any = None

        try:
            import boto3

            self._client = boto3.client("autoscaling", region_name=region)
            self._aws_available = True
        except ImportError:
            logger.warning("boto3 not installed, AWS backend disabled")
        except Exception as e:
            logger.warning(f"Failed to initialize AWS client: {e}")

    def get_worker_count(self) -> int:
        """Get current desired capacity from ASG."""
        if not self._aws_available:
            logger.warning("AWS not available, returning 0")
            return 0

        try:
            response = self._client.describe_auto_scaling_groups(
                AutoScalingGroupNames=[self.asg_name]
            )
            groups = response.get("AutoScalingGroups", [])
            if groups:
                return groups[0].get("DesiredCapacity", 0)
            return 0
        except Exception as e:
            logger.error(f"Failed to get worker count: {e}")
            return 0

    def scale_up(self, count: int) -> bool:
        """Scale up by increasing desired capacity."""
        if not self._aws_available:
            logger.error("AWS not available")
            return False

        try:
            current = self.get_worker_count()
            return self.set_worker_count(current + count)
        except Exception as e:
            logger.error(f"Failed to scale up: {e}")
            return False

    def scale_down(self, count: int) -> bool:
        """Scale down by decreasing desired capacity."""
        if not self._aws_available:
            logger.error("AWS not available")
            return False

        try:
            current = self.get_worker_count()
            return self.set_worker_count(max(0, current - count))
        except Exception as e:
            logger.error(f"Failed to scale down: {e}")
            return False

    def set_worker_count(self, count: int) -> bool:
        """Set exact desired capacity on the ASG."""
        if not self._aws_available:
            logger.error("AWS not available")
            return False

        try:
            self._client.set_desired_capacity(
                AutoScalingGroupName=self.asg_name,
                DesiredCapacity=count,
                HonorCooldown=False,
            )
            logger.info(f"AWS ASG: set {self.asg_name} desired capacity to {count}")
            return True
        except Exception as e:
            logger.error(f"Failed to set worker count: {e}")
            return False


class GcpMigBackend:
    """GCP Managed Instance Group backend.

    Requires google-cloud-compute to be installed.
    """

    def __init__(
        self,
        project_id: str,
        zone: str,
        instance_group_name: str,
    ) -> None:
        """Initialize GCP MIG backend.

        Args:
            project_id: GCP project ID
            zone: GCP zone (e.g., 'us-central1-a')
            instance_group_name: Name of the managed instance group
        """
        self.project_id = project_id
        self.zone = zone
        self.instance_group_name = instance_group_name
        self._gcp_available = False
        self._client: Any = None

        try:
            from google.cloud import compute_v1

            self._client = compute_v1.InstanceGroupManagersClient()
            self._gcp_available = True
        except ImportError:
            logger.warning("google-cloud-compute not installed, GCP backend disabled")
        except Exception as e:
            logger.warning(f"Failed to initialize GCP client: {e}")

    def get_worker_count(self) -> int:
        """Get current target size from MIG."""
        if not self._gcp_available:
            logger.warning("GCP not available, returning 0")
            return 0

        try:
            mig = self._client.get(
                project=self.project_id,
                zone=self.zone,
                instance_group_manager=self.instance_group_name,
            )
            return mig.target_size or 0
        except Exception as e:
            logger.error(f"Failed to get worker count: {e}")
            return 0

    def scale_up(self, count: int) -> bool:
        """Scale up by increasing MIG target size."""
        if not self._gcp_available:
            logger.error("GCP not available")
            return False

        try:
            current = self.get_worker_count()
            return self.set_worker_count(current + count)
        except Exception as e:
            logger.error(f"Failed to scale up: {e}")
            return False

    def scale_down(self, count: int) -> bool:
        """Scale down by decreasing MIG target size."""
        if not self._gcp_available:
            logger.error("GCP not available")
            return False

        try:
            current = self.get_worker_count()
            return self.set_worker_count(max(0, current - count))
        except Exception as e:
            logger.error(f"Failed to scale down: {e}")
            return False

    def set_worker_count(self, count: int) -> bool:
        """Set exact target size on the MIG."""
        if not self._gcp_available:
            logger.error("GCP not available")
            return False

        try:
            # Resize the MIG
            operation = self._client.resize(
                project=self.project_id,
                zone=self.zone,
                instance_group_manager=self.instance_group_name,
                size=count,
            )
            logger.info(
                f"GCP MIG: set {self.instance_group_name} target size to {count} "
                f"(operation: {operation.name})"
            )
            return True
        except Exception as e:
            logger.error(f"Failed to set worker count: {e}")
            return False


class AutoScaler:
    """Auto-scaler for worker pools based on queue metrics.

    Monitors queue depth and processing rate to make scaling decisions.
    Supports multiple backends for different deployment environments.

    Example:
        >>> from oclawma.autoscale import AutoScaler, ScalingThresholds
        >>> thresholds = ScalingThresholds(
        ...     scale_up_threshold=50,
        ...     scale_down_threshold=5,
        ...     max_workers=20,
        ...     min_workers=2,
        ... )
        >>> scaler = AutoScaler(thresholds=thresholds)
        >>> recommendation = scaler.recommend_scale(queue_depth=100, worker_count=5)
        >>> if recommendation.decision == ScalingDecision.SCALE_UP:
        ...     scaler.scale_up(recommendation.target_workers - recommendation.current_workers)
    """

    def __init__(
        self,
        thresholds: ScalingThresholds | None = None,
        backend: WorkerPoolBackend | None = None,
        metrics_collector: Any | None = None,
    ) -> None:
        """Initialize the auto-scaler.

        Args:
            thresholds: Scaling thresholds configuration
            backend: Worker pool backend (defaults to LocalWorkerBackend)
            metrics_collector: Optional metrics collector for integration
        """
        self.thresholds = thresholds or ScalingThresholds()
        self.backend = backend or LocalWorkerBackend()
        self.metrics_collector = metrics_collector

        self._lock = threading.RLock()
        self._last_scale_up_time: float = 0.0
        self._last_scale_down_time: float = 0.0
        self._processing_history: list[tuple[float, int]] = []  # (timestamp, jobs_processed)
        self._current_worker_count: int = 0

        # Get initial worker count from backend
        self._sync_worker_count()

    def _sync_worker_count(self) -> None:
        """Synchronize worker count from backend."""
        with self._lock:
            self._current_worker_count = self.backend.get_worker_count()

    def update_processing_count(self, jobs_processed: int) -> None:
        """Update the processing count for rate calculation.

        Args:
            jobs_processed: Total jobs processed (monotonically increasing)
        """
        with self._lock:
            now = time.time()
            self._processing_history.append((now, jobs_processed))

            # Prune old entries outside the window
            cutoff = now - self.thresholds.processing_rate_window
            self._processing_history = [(t, c) for t, c in self._processing_history if t >= cutoff]

    def get_processing_rate(self) -> float:
        """Calculate the current processing rate (jobs per second).

        Returns:
            Jobs processed per second over the configured window
        """
        with self._lock:
            if len(self._processing_history) < 2:
                return 0.0

            # Sort by timestamp
            sorted_history = sorted(self._processing_history, key=lambda x: x[0])
            first_time, first_count = sorted_history[0]
            last_time, last_count = sorted_history[-1]

            time_delta = last_time - first_time
            if time_delta <= 0:
                return 0.0

            count_delta = last_count - first_count
            return max(0.0, count_delta / time_delta)

    def recommend_scale(
        self,
        queue_depth: int,
        worker_count: int | None = None,
    ) -> ScalingRecommendation:
        """Recommend scaling action based on current metrics.

        Args:
            queue_depth: Current number of jobs in queue
            worker_count: Current worker count (uses backend if not provided)

        Returns:
            Scaling recommendation with target worker count and reasoning
        """
        with self._lock:
            if worker_count is None:
                worker_count = self._current_worker_count

            processing_rate = self.get_processing_rate()
            metrics = ScalingMetrics(
                queue_depth=queue_depth,
                processing_rate=processing_rate,
                worker_count=worker_count,
            )

            now = time.time()

            # Check cooldowns
            scale_up_cooldown_remaining = max(
                0, self.thresholds.scale_up_cooldown - (now - self._last_scale_up_time)
            )
            scale_down_cooldown_remaining = max(
                0, self.thresholds.scale_down_cooldown - (now - self._last_scale_down_time)
            )

            # Calculate optimal workers based on queue depth
            # Simple heuristic: 1 worker per 10 jobs in queue
            jobs_per_worker = max(1, self.thresholds.scale_up_threshold // max(1, worker_count))
            optimal_workers = max(
                self.thresholds.min_workers,
                min(
                    self.thresholds.max_workers,
                    (queue_depth + jobs_per_worker - 1) // jobs_per_worker,
                ),
            )

            # Decision logic
            if queue_depth >= self.thresholds.scale_up_threshold:
                if scale_up_cooldown_remaining > 0:
                    return ScalingRecommendation(
                        decision=ScalingDecision.MAINTAIN,
                        target_workers=worker_count,
                        current_workers=worker_count,
                        reason=f"Scale-up needed but in cooldown ({scale_up_cooldown_remaining:.0f}s remaining)",
                        metrics=metrics,
                    )

                # Scale up to handle queue
                target = min(optimal_workers + 1, self.thresholds.max_workers)
                if target > worker_count:
                    return ScalingRecommendation(
                        decision=ScalingDecision.SCALE_UP,
                        target_workers=target,
                        current_workers=worker_count,
                        reason=f"Queue depth ({queue_depth}) exceeds threshold ({self.thresholds.scale_up_threshold})",
                        metrics=metrics,
                    )

            elif queue_depth <= self.thresholds.scale_down_threshold:
                if scale_down_cooldown_remaining > 0:
                    return ScalingRecommendation(
                        decision=ScalingDecision.MAINTAIN,
                        target_workers=worker_count,
                        current_workers=worker_count,
                        reason=f"Scale-down possible but in cooldown ({scale_down_cooldown_remaining:.0f}s remaining)",
                        metrics=metrics,
                    )

                # Scale down if over-provisioned
                target = max(optimal_workers - 1, self.thresholds.min_workers)
                if target < worker_count:
                    return ScalingRecommendation(
                        decision=ScalingDecision.SCALE_DOWN,
                        target_workers=target,
                        current_workers=worker_count,
                        reason=f"Queue depth ({queue_depth}) below threshold ({self.thresholds.scale_down_threshold})",
                        metrics=metrics,
                    )

            # Maintain current scale
            return ScalingRecommendation(
                decision=ScalingDecision.MAINTAIN,
                target_workers=worker_count,
                current_workers=worker_count,
                reason=f"Queue depth ({queue_depth}) within normal range",
                metrics=metrics,
            )

    def scale_up(self, count: int) -> bool:
        """Scale up the worker pool.

        Args:
            count: Number of workers to add

        Returns:
            True if scaling was successful
        """
        with self._lock:
            current = self.backend.get_worker_count()
            if current + count > self.thresholds.max_workers:
                count = self.thresholds.max_workers - current

            if count <= 0:
                logger.debug("Scale up: already at max workers")
                return False

            success = self.backend.scale_up(count)
            if success:
                self._last_scale_up_time = time.time()
                self._current_worker_count = self.backend.get_worker_count()

                # Update metrics collector if provided
                if self.metrics_collector:
                    with suppress(Exception):
                        self.metrics_collector.update_job_status_count(
                            "workers", self._current_worker_count
                        )

            return success

    def scale_down(self, count: int) -> bool:
        """Scale down the worker pool.

        Args:
            count: Number of workers to remove

        Returns:
            True if scaling was successful
        """
        with self._lock:
            current = self.backend.get_worker_count()
            if current - count < self.thresholds.min_workers:
                count = current - self.thresholds.min_workers

            if count <= 0:
                logger.debug("Scale down: already at min workers")
                return False

            success = self.backend.scale_down(count)
            if success:
                self._last_scale_down_time = time.time()
                self._current_worker_count = self.backend.get_worker_count()

                # Update metrics collector if provided
                if self.metrics_collector:
                    with suppress(Exception):
                        self.metrics_collector.update_job_status_count(
                            "workers", self._current_worker_count
                        )

            return success

    def apply_recommendation(self, recommendation: ScalingRecommendation) -> bool:
        """Apply a scaling recommendation.

        Args:
            recommendation: Scaling recommendation to apply

        Returns:
            True if any scaling action was taken
        """
        if recommendation.decision == ScalingDecision.SCALE_UP:
            delta = recommendation.target_workers - recommendation.current_workers
            return self.scale_up(delta) if delta > 0 else False
        elif recommendation.decision == ScalingDecision.SCALE_DOWN:
            delta = recommendation.current_workers - recommendation.target_workers
            return self.scale_down(delta) if delta > 0 else False
        else:
            return False

    def set_worker_count(self, count: int) -> bool:
        """Set exact worker count.

        Args:
            count: Target number of workers

        Returns:
            True if successful
        """
        with self._lock:
            success = self.backend.set_worker_count(count)
            if success:
                self._current_worker_count = self.backend.get_worker_count()
            return success

    def get_current_metrics(self) -> ScalingMetrics:
        """Get current scaling metrics.

        Returns:
            Current metrics snapshot
        """
        with self._lock:
            return ScalingMetrics(
                queue_depth=self.metrics_collector.queue_depth if self.metrics_collector else 0,
                processing_rate=self.get_processing_rate(),
                worker_count=self._current_worker_count,
            )

    def get_status(self) -> dict[str, Any]:
        """Get complete auto-scaler status.

        Returns:
            Dictionary with current status and configuration
        """
        with self._lock:
            now = time.time()
            return {
                "current_workers": self._current_worker_count,
                "thresholds": {
                    "scale_up_threshold": self.thresholds.scale_up_threshold,
                    "scale_down_threshold": self.thresholds.scale_down_threshold,
                    "max_workers": self.thresholds.max_workers,
                    "min_workers": self.thresholds.min_workers,
                    "scale_up_cooldown": self.thresholds.scale_up_cooldown,
                    "scale_down_cooldown": self.thresholds.scale_down_cooldown,
                },
                "cooldowns": {
                    "scale_up_remaining": max(
                        0, self.thresholds.scale_up_cooldown - (now - self._last_scale_up_time)
                    ),
                    "scale_down_remaining": max(
                        0, self.thresholds.scale_down_cooldown - (now - self._last_scale_down_time)
                    ),
                },
                "processing_rate": self.get_processing_rate(),
                "backend_type": type(self.backend).__name__,
            }


def create_autoscaler(
    backend_type: str = "local",
    thresholds: ScalingThresholds | None = None,
    metrics_collector: Any | None = None,
    **backend_kwargs: Any,
) -> AutoScaler:
    """Factory function to create an auto-scaler with specified backend.

    Args:
        backend_type: Type of backend ('local', 'kubernetes', 'aws', 'gcp')
        thresholds: Scaling thresholds configuration
        metrics_collector: Optional metrics collector for integration
        **backend_kwargs: Backend-specific configuration

    Returns:
        Configured AutoScaler instance

    Raises:
        ValueError: If backend_type is not recognized
    """
    backend: WorkerPoolBackend

    if backend_type == "local":
        backend = LocalWorkerBackend(initial_workers=backend_kwargs.get("initial_workers", 1))
    elif backend_type == "kubernetes":
        backend = KubernetesHpaBackend(
            deployment_name=backend_kwargs.get("deployment_name", "oclawma-workers"),
            namespace=backend_kwargs.get("namespace", "default"),
            hpa_name=backend_kwargs.get("hpa_name"),
        )
    elif backend_type == "aws":
        backend = AwsAsgBackend(
            asg_name=backend_kwargs.get("asg_name", "oclawma-asg"),
            region=backend_kwargs.get("region", "us-east-1"),
        )
    elif backend_type == "gcp":
        backend = GcpMigBackend(
            project_id=backend_kwargs.get("project_id", ""),
            zone=backend_kwargs.get("zone", "us-central1-a"),
            instance_group_name=backend_kwargs.get("instance_group_name", "oclawma-mig"),
        )
    else:
        raise ValueError(f"Unknown backend type: {backend_type}")

    return AutoScaler(
        thresholds=thresholds,
        backend=backend,
        metrics_collector=metrics_collector,
    )
