from __future__ import annotations

from collections import defaultdict, deque
import threading
import time
from dataclasses import dataclass, replace
import math
from typing import Callable, Iterable, Any

from flownet.api.scheduler import SchedulePlan, SimpleScheduler
from flownet.runtime.cluster_view import ClusterNode, ClusterView


DEFAULT_LOCALITY_ORDER = ("local", "zone", "region", "remote")
DEFAULT_RESERVATION_TIMEOUT_MS = 30000
DEFAULT_LATENCY_SAMPLE_WINDOW = 2048


@dataclass(frozen=True)
class PlanResolutionRequest:
    stage_id: str
    request_id: str
    num_partitions: int
    affinity_hint: str | None = None
    addresses: tuple[str, ...] | None = None
    merge_addresses: tuple[str, ...] | None = None
    merge_partitions: dict[str, str] | None = None
    include_merge_address: bool = True
    merge_sharded: bool = False
    epoch: int = 1
    plan_id: str | None = None
    require_healthy: bool = True
    locality_order: tuple[str, ...] = DEFAULT_LOCALITY_ORDER
    resource_requirements: dict[str, int | float] | None = None
    allow_resource_fallback: bool = True
    lease_owner_kind: str | None = None
    lease_owner_id: str | None = None
    reservation_timeout_ms: int | None = None
    autoscale_min_workers: int | None = None
    autoscale_max_workers: int | None = None
    autoscale_queue_depth: int | None = None
    autoscale_target_queue_per_worker: int | None = None
    autoscale_observed_latency_ms: float | None = None
    autoscale_latency_budget_ms: float | None = None
    autoscale_observed_throughput_rps: float | None = None
    autoscale_target_throughput_rps: float | None = None
    autoscale_cooldown_ms: int | None = None
    spillover_stage1_tiers: tuple[str, ...] = ("local", "zone")
    spillover_stage2_tiers: tuple[str, ...] = ("region",)
    spillover_allow_remote: bool = True
    spillover_max_remote_nodes: int | None = None
    spillover_ttl_ms: int | None = None
    spillover_claim_id: str | None = None


@dataclass(frozen=True)
class PlanResolution:
    plan: SchedulePlan
    addresses: tuple[str, ...]
    merge_addresses: tuple[str, ...]
    strategy: str
    resource_fallback_used: bool = False


@dataclass
class _AutoscaleStageState:
    desired_workers: int
    cooldown_ms: int
    last_scale_ts: float = 0.0
    last_direction: str = "steady"
    last_reason: str = "steady"
    last_decision_ts: float = 0.0
    decisions: int = 0
    cooldown_suppressed: int = 0


class ClusterPlanResolver:
    """Resolve scheduler plans from ClusterView using locality + health ranking.

    Design guardrail:
    - Resolver consumes ClusterView snapshots and emits placement plans.
    - Resolver must not mutate/own membership lifecycle; ClusterView maintenance
      stays in runtime cluster/gossip paths.
    """

    def __init__(
        self,
        cluster_view: ClusterView,
        *,
        time_fn: Callable[[], float] | None = None,
    ):
        self._cluster_view = cluster_view
        self._time_fn = time_fn or time.time
        self._metrics_lock = threading.Lock()
        self._plan_resolution_count = 0
        self._resource_fallback_count = 0
        self._locality_hits: dict[str, int] = {
            "local": 0,
            "zone": 0,
            "region": 0,
            "remote": 0,
            "unknown": 0,
        }
        self._locality_total_hits = 0
        self._staleness_count = 0
        self._staleness_ms_sum = 0.0
        self._staleness_ms_max = 0.0
        self._staleness_ms_latest = 0.0
        self._latency_count = 0
        self._latency_ms_sum = 0.0
        self._latency_ms_min: float | None = None
        self._latency_ms_max = 0.0
        self._latency_ms_recent: deque[float] = deque(maxlen=DEFAULT_LATENCY_SAMPLE_WINDOW)
        # Owner-aware fairness state (best-effort in-memory lease reservations).
        self._owner_resolution_count = 0
        self._reservation_entries: list[tuple[str, str, float]] = []
        self._owner_active_reservations: dict[str, int] = defaultdict(int)
        self._node_active_reservations: dict[str, int] = defaultdict(int)
        self._owner_node_active_reservations: dict[tuple[str, str], int] = defaultdict(int)
        # Autoscaling control-loop state (best-effort, stage-local).
        self._autoscale_states: dict[str, _AutoscaleStageState] = {}
        self._autoscale_decision_count = 0
        self._autoscale_scale_up_count = 0
        self._autoscale_scale_down_count = 0
        self._autoscale_cooldown_suppressed_count = 0
        # Spillover control state (best-effort, in-memory).
        self._spillover_decision_count = 0
        self._spillover_remote_used_count = 0
        self._spillover_remote_blocked_count = 0
        self._spillover_claim_created_count = 0
        self._spillover_claim_reused_count = 0
        self._spillover_claim_expiry: dict[str, float] = {}
        # Request idempotence for owner lease reservations.
        self._owner_request_claim_expiry: dict[str, float] = {}

    def resolve(self, request: PlanResolutionRequest) -> PlanResolution:
        started_at = self._time_fn()
        self._prune_expired_reservations(now=started_at)
        effective_request, autoscale_applied = self._apply_autoscale(
            request=request,
            now=started_at,
        )
        addresses, resource_fallback_used, spillover_strategy = self._resolve_addresses(
            effective_request,
            now=started_at,
        )
        merge_addresses = tuple(effective_request.merge_addresses or addresses)
        plan = SimpleScheduler(addresses).make_plan(
            stage_id=effective_request.stage_id,
            request_id=effective_request.request_id,
            num_partitions=effective_request.num_partitions,
            merge_addresses=merge_addresses,
            merge_partitions=effective_request.merge_partitions,
            include_merge_address=effective_request.include_merge_address,
            merge_sharded=effective_request.merge_sharded,
            epoch=effective_request.epoch,
            plan_id=effective_request.plan_id,
        )
        strategy = "explicit_addresses" if request.addresses else "cluster_view_locality_health"
        if resource_fallback_used:
            strategy = f"{strategy}+resource_fallback"
        if autoscale_applied:
            strategy = f"{strategy}+autoscale"
        if spillover_strategy:
            strategy = f"{strategy}+{spillover_strategy}"
        completed_at = self._time_fn()
        self._record_owner_reservations(
            request=effective_request,
            plan=plan,
            now=completed_at,
        )
        self._record_metrics(
            latency_ms=(completed_at - started_at) * 1000.0,
            staleness_ms=self._cluster_view.estimate_staleness_ms(now=completed_at),
            resource_fallback_used=resource_fallback_used,
            locality_hits=self._count_locality_hits(plan),
        )
        return PlanResolution(
            plan=plan,
            addresses=addresses,
            merge_addresses=merge_addresses,
            strategy=strategy,
            resource_fallback_used=resource_fallback_used,
        )

    def resolve_plan(
        self,
        *,
        stage_id: str,
        request_id: str,
        num_partitions: int,
        affinity_hint: str | None = None,
        addresses: Iterable[str] | None = None,
        merge_addresses: Iterable[str] | None = None,
        merge_partitions: dict[str, str] | None = None,
        include_merge_address: bool = True,
        merge_sharded: bool = False,
        epoch: int = 1,
        plan_id: str | None = None,
        require_healthy: bool = True,
        locality_order: tuple[str, ...] = DEFAULT_LOCALITY_ORDER,
        resource_requirements: dict[str, int | float] | None = None,
        allow_resource_fallback: bool = True,
        lease_owner_kind: str | None = None,
        lease_owner_id: str | None = None,
        reservation_timeout_ms: int | None = None,
        autoscale_min_workers: int | None = None,
        autoscale_max_workers: int | None = None,
        autoscale_queue_depth: int | None = None,
        autoscale_target_queue_per_worker: int | None = None,
        autoscale_observed_latency_ms: float | None = None,
        autoscale_latency_budget_ms: float | None = None,
        autoscale_observed_throughput_rps: float | None = None,
        autoscale_target_throughput_rps: float | None = None,
        autoscale_cooldown_ms: int | None = None,
        spillover_stage1_tiers: tuple[str, ...] = ("local", "zone"),
        spillover_stage2_tiers: tuple[str, ...] = ("region",),
        spillover_allow_remote: bool = True,
        spillover_max_remote_nodes: int | None = None,
        spillover_ttl_ms: int | None = None,
        spillover_claim_id: str | None = None,
    ) -> SchedulePlan:
        request = PlanResolutionRequest(
            stage_id=stage_id,
            request_id=request_id,
            num_partitions=num_partitions,
            affinity_hint=affinity_hint,
            addresses=tuple(addresses) if addresses is not None else None,
            merge_addresses=tuple(merge_addresses) if merge_addresses is not None else None,
            merge_partitions=merge_partitions,
            include_merge_address=include_merge_address,
            merge_sharded=merge_sharded,
            epoch=epoch,
            plan_id=plan_id,
            require_healthy=require_healthy,
            locality_order=locality_order,
            resource_requirements=resource_requirements,
            allow_resource_fallback=allow_resource_fallback,
            lease_owner_kind=lease_owner_kind,
            lease_owner_id=lease_owner_id,
            reservation_timeout_ms=reservation_timeout_ms,
            autoscale_min_workers=autoscale_min_workers,
            autoscale_max_workers=autoscale_max_workers,
            autoscale_queue_depth=autoscale_queue_depth,
            autoscale_target_queue_per_worker=autoscale_target_queue_per_worker,
            autoscale_observed_latency_ms=autoscale_observed_latency_ms,
            autoscale_latency_budget_ms=autoscale_latency_budget_ms,
            autoscale_observed_throughput_rps=autoscale_observed_throughput_rps,
            autoscale_target_throughput_rps=autoscale_target_throughput_rps,
            autoscale_cooldown_ms=autoscale_cooldown_ms,
            spillover_stage1_tiers=spillover_stage1_tiers,
            spillover_stage2_tiers=spillover_stage2_tiers,
            spillover_allow_remote=spillover_allow_remote,
            spillover_max_remote_nodes=spillover_max_remote_nodes,
            spillover_ttl_ms=spillover_ttl_ms,
            spillover_claim_id=spillover_claim_id,
        )
        return self.resolve(request).plan

    def metrics_snapshot(self) -> dict[str, object]:
        now = self._time_fn()
        with self._metrics_lock:
            self._prune_expired_locked(now=now)
            total = self._plan_resolution_count
            fallback_rate = (
                float(self._resource_fallback_count) / float(total) if total else 0.0
            )
            hit_total = self._locality_total_hits
            locality_rate = {
                tier: (float(count) / float(hit_total) if hit_total else 0.0)
                for tier, count in self._locality_hits.items()
            }
            avg_staleness = (
                self._staleness_ms_sum / float(self._staleness_count)
                if self._staleness_count
                else 0.0
            )
            avg_latency = (
                self._latency_ms_sum / float(self._latency_count)
                if self._latency_count
                else 0.0
            )
            owner_active = dict(self._owner_active_reservations)
            owner_counts = list(owner_active.values())
            max_owner_active = max(owner_counts) if owner_counts else 0
            min_owner_active = min(owner_counts) if owner_counts else 0
            return {
                "cluster_view_staleness_ms": {
                    "latest": self._staleness_ms_latest,
                    "avg": avg_staleness,
                    "max": self._staleness_ms_max,
                    "samples": self._staleness_count,
                },
                "scheduler_resource_fallback_rate": {
                    "rate": fallback_rate,
                    "fallback_count": self._resource_fallback_count,
                    "total": total,
                },
                "scheduler_locality_tier_hit_rate": {
                    "rate": locality_rate,
                    "counts": dict(self._locality_hits),
                    "total_hits": hit_total,
                },
                "plan_resolution_latency_ms": {
                    "avg": avg_latency,
                    "min": 0.0 if self._latency_ms_min is None else self._latency_ms_min,
                    "max": self._latency_ms_max,
                    "samples": self._latency_count,
                    "p50": _quantile(self._latency_ms_recent, 0.50),
                    "p95": _quantile(self._latency_ms_recent, 0.95),
                    "quantile_window_samples": len(self._latency_ms_recent),
                },
                "scheduler_owner_fairness": {
                    "owner_resolution_count": self._owner_resolution_count,
                    "active_owners": len(owner_active),
                    "active_reservations": len(self._reservation_entries),
                    "max_owner_active_reservations": max_owner_active,
                    "min_owner_active_reservations": min_owner_active,
                    "owner_active_reservations": owner_active,
                    "node_active_reservations": dict(self._node_active_reservations),
                },
                "scheduler_spillover": {
                    "decision_count": self._spillover_decision_count,
                    "remote_used_count": self._spillover_remote_used_count,
                    "remote_blocked_count": self._spillover_remote_blocked_count,
                    "claim_created_count": self._spillover_claim_created_count,
                    "claim_reused_count": self._spillover_claim_reused_count,
                    "active_claims": len(self._spillover_claim_expiry),
                },
                "scheduler_autoscaling": self._autoscale_snapshot_locked(),
            }

    def reset_metrics(self) -> None:
        with self._metrics_lock:
            self._plan_resolution_count = 0
            self._resource_fallback_count = 0
            for key in list(self._locality_hits.keys()):
                self._locality_hits[key] = 0
            self._locality_total_hits = 0
            self._staleness_count = 0
            self._staleness_ms_sum = 0.0
            self._staleness_ms_max = 0.0
            self._staleness_ms_latest = 0.0
            self._latency_count = 0
            self._latency_ms_sum = 0.0
            self._latency_ms_min = None
            self._latency_ms_max = 0.0
            self._latency_ms_recent.clear()
            self._owner_resolution_count = 0
            self._reservation_entries = []
            self._owner_active_reservations.clear()
            self._node_active_reservations.clear()
            self._owner_node_active_reservations.clear()
            self._autoscale_states.clear()
            self._autoscale_decision_count = 0
            self._autoscale_scale_up_count = 0
            self._autoscale_scale_down_count = 0
            self._autoscale_cooldown_suppressed_count = 0
            self._spillover_decision_count = 0
            self._spillover_remote_used_count = 0
            self._spillover_remote_blocked_count = 0
            self._spillover_claim_created_count = 0
            self._spillover_claim_reused_count = 0
            self._spillover_claim_expiry.clear()
            self._owner_request_claim_expiry.clear()

    def _resolve_addresses(
        self,
        request: PlanResolutionRequest,
        *,
        now: float,
    ) -> tuple[tuple[str, ...], bool, str | None]:
        if request.addresses:
            addresses = _dedupe_preserve_order(request.addresses)
            if addresses:
                return addresses, False, None
            raise ValueError("PlanResolutionRequest.addresses must not be empty.")

        ranked, resource_fallback_used, spillover_strategy = self._ranked_nodes(request, now=now)
        if request.affinity_hint:
            hint_node = self._cluster_view.resolve_node_for_key(
                request.stage_id,
                affinity_hint=request.affinity_hint,
            )
            ranked = _promote_node(ranked, hint_node.node_id)
        addresses = _dedupe_preserve_order(node.address for node in ranked)
        if not addresses:
            raise ValueError("No schedulable nodes available from ClusterView.")
        return addresses, resource_fallback_used, spillover_strategy

    def _ranked_nodes(
        self,
        request: PlanResolutionRequest,
        *,
        now: float,
    ) -> tuple[list[ClusterNode], bool, str | None]:
        nodes = [
            node
            for node in self._cluster_view.list_nodes()
            if self._cluster_view.is_node_schedulable(node.node_id)
        ]
        if request.require_healthy:
            healthy_nodes = [
                node for node in nodes if self._cluster_view.is_node_healthy(node.node_id)
            ]
            if healthy_nodes:
                nodes = healthy_nodes
        resource_fallback_used = False
        if request.resource_requirements:
            matched = [
                node
                for node in nodes
                if _resource_match(node.resource_summary, request.resource_requirements)
            ]
            if matched:
                nodes = matched
            elif not request.allow_resource_fallback:
                raise ValueError(
                    "No nodes satisfy resource_requirements and allow_resource_fallback=False."
                )
            else:
                resource_fallback_used = True

        order = {tier: idx for idx, tier in enumerate(request.locality_order)}
        max_priority = len(order)
        owner_key = _owner_key(request)
        owner_node_counts: dict[str, int] = {}
        node_counts: dict[str, int] = {}
        if owner_key is not None:
            with self._metrics_lock:
                owner_node_counts = {
                    node_id: count
                    for (entry_owner, node_id), count in self._owner_node_active_reservations.items()
                    if entry_owner == owner_key
                }
                node_counts = dict(self._node_active_reservations)

        def sort_key(node: ClusterNode):
            tier = self._cluster_view.locality_tier(node.node_id)
            priority = order.get(tier, max_priority)
            freshness = 0.0 if node.last_seen is None else node.last_seen
            if owner_key is None:
                return (priority, -freshness, node.node_id)
            owner_load = owner_node_counts.get(node.node_id, 0)
            global_load = node_counts.get(node.node_id, 0)
            return (priority, owner_load, global_load, -freshness, node.node_id)

        ranked = sorted(nodes, key=sort_key)
        ranked, spillover_strategy = self._apply_spillover_policy(
            ranked=ranked,
            request=request,
            now=now,
        )
        return ranked, resource_fallback_used, spillover_strategy

    def _record_owner_reservations(
        self,
        *,
        request: PlanResolutionRequest,
        plan: SchedulePlan,
        now: float,
    ) -> None:
        owner_key = _owner_key(request)
        if owner_key is None:
            return

        ttl_ms = request.reservation_timeout_ms
        if ttl_ms is None:
            ttl_ms = DEFAULT_RESERVATION_TIMEOUT_MS
        try:
            ttl_ms_int = int(ttl_ms)
        except (TypeError, ValueError) as exc:
            raise ValueError("reservation_timeout_ms must be an integer > 0.") from exc
        if ttl_ms_int <= 0:
            raise ValueError("reservation_timeout_ms must be an integer > 0.")

        expires_at = now + (float(ttl_ms_int) / 1000.0)
        node_id_by_address = {node.address: node.node_id for node in self._cluster_view.list_nodes()}

        with self._metrics_lock:
            self._prune_expired_locked(now=now)
            request_key = f"{owner_key}:{request.stage_id}:{request.request_id}"
            prev_claim_expiry = float(self._owner_request_claim_expiry.get(request_key, 0.0))
            if prev_claim_expiry > now:
                return
            self._owner_resolution_count += 1
            for address in plan.partitions.values():
                node_id = node_id_by_address.get(address, f"addr:{address}")
                self._reservation_entries.append((owner_key, node_id, expires_at))
                self._owner_active_reservations[owner_key] += 1
                self._node_active_reservations[node_id] += 1
                self._owner_node_active_reservations[(owner_key, node_id)] += 1
            self._owner_request_claim_expiry[request_key] = expires_at

    def _prune_expired_reservations(self, *, now: float) -> None:
        with self._metrics_lock:
            self._prune_expired_locked(now=now)

    def _prune_expired_locked(self, *, now: float) -> None:
        expired_spillover_claims = [
            key
            for key, expires_at in self._spillover_claim_expiry.items()
            if expires_at <= now
        ]
        for key in expired_spillover_claims:
            self._spillover_claim_expiry.pop(key, None)

        expired_owner_claims = [
            key
            for key, expires_at in self._owner_request_claim_expiry.items()
            if expires_at <= now
        ]
        for key in expired_owner_claims:
            self._owner_request_claim_expiry.pop(key, None)

        if not self._reservation_entries:
            return

        alive_entries: list[tuple[str, str, float]] = []
        owner_active: dict[str, int] = defaultdict(int)
        node_active: dict[str, int] = defaultdict(int)
        owner_node_active: dict[tuple[str, str], int] = defaultdict(int)

        for owner_key, node_id, expires_at in self._reservation_entries:
            if expires_at <= now:
                continue
            alive_entries.append((owner_key, node_id, expires_at))
            owner_active[owner_key] += 1
            node_active[node_id] += 1
            owner_node_active[(owner_key, node_id)] += 1

        self._reservation_entries = alive_entries
        self._owner_active_reservations = owner_active
        self._node_active_reservations = node_active
        self._owner_node_active_reservations = owner_node_active

    def _apply_spillover_policy(
        self,
        *,
        ranked: list[ClusterNode],
        request: PlanResolutionRequest,
        now: float,
    ) -> tuple[list[ClusterNode], str | None]:
        if not ranked:
            return ranked, None
        if not _spillover_policy_enabled(request):
            return ranked, None

        target_nodes = min(max(1, int(request.num_partitions)), len(ranked))
        stage1 = set(request.spillover_stage1_tiers)
        stage2 = set(request.spillover_stage2_tiers)
        selected: list[ClusterNode] = []
        selected_ids: set[str] = set()

        def _append_from_tiers(tiers: set[str]) -> None:
            if len(selected) >= target_nodes:
                return
            for node in ranked:
                if len(selected) >= target_nodes:
                    return
                if node.node_id in selected_ids:
                    continue
                tier = self._cluster_view.locality_tier(node.node_id)
                if tier not in tiers:
                    continue
                selected.append(node)
                selected_ids.add(node.node_id)

        _append_from_tiers(stage1)
        _append_from_tiers(stage2)

        strategy_parts: list[str] = []
        remote_needed = len(selected) < target_nodes
        remote_used = False
        remote_blocked = False
        claim_created = False
        claim_reused = False
        if remote_needed:
            remote_pool = [
                node
                for node in ranked
                if node.node_id not in selected_ids
                and self._cluster_view.locality_tier(node.node_id) not in stage1.union(stage2)
            ]
            allow_remote = request.spillover_allow_remote
            if allow_remote and request.spillover_ttl_ms is not None:
                allow_remote, claim_created, claim_reused = self._acquire_spillover_claim(
                    request=request,
                    now=now,
                )
            if allow_remote:
                max_remote = request.spillover_max_remote_nodes
                if max_remote is not None:
                    try:
                        max_remote = max(0, int(max_remote))
                    except (TypeError, ValueError):
                        max_remote = 0
                remaining = target_nodes - len(selected)
                allowed = len(remote_pool)
                if max_remote is not None:
                    allowed = min(allowed, max_remote)
                allowed = min(allowed, remaining)
                for node in remote_pool[:allowed]:
                    selected.append(node)
                    selected_ids.add(node.node_id)
                remote_used = allowed > 0
                remote_blocked = remaining > allowed
            else:
                remote_blocked = True

        if not selected:
            return ranked, None

        if remote_used:
            strategy_parts.append("spillover_remote")
        if remote_blocked:
            strategy_parts.append("spillover_limited")
        with self._metrics_lock:
            self._prune_expired_locked(now=now)
            self._spillover_decision_count += 1
            if remote_used:
                self._spillover_remote_used_count += 1
            if remote_blocked:
                self._spillover_remote_blocked_count += 1
            if claim_created:
                self._spillover_claim_created_count += 1
            if claim_reused:
                self._spillover_claim_reused_count += 1
        return selected, "+".join(strategy_parts) if strategy_parts else "spillover_local_first"

    def _acquire_spillover_claim(
        self,
        *,
        request: PlanResolutionRequest,
        now: float,
    ) -> tuple[bool, bool, bool]:
        ttl_ms = request.spillover_ttl_ms
        if ttl_ms is None:
            return True, False, False
        try:
            ttl_ms_int = int(ttl_ms)
        except (TypeError, ValueError):
            return False, False, False
        if ttl_ms_int <= 0:
            return False, False, False

        claim_key = _spillover_claim_key(request)
        expires_at = now + (float(ttl_ms_int) / 1000.0)
        with self._metrics_lock:
            self._prune_expired_locked(now=now)
            prev = float(self._spillover_claim_expiry.get(claim_key, 0.0))
            if prev > now:
                return True, False, True
            self._spillover_claim_expiry[claim_key] = expires_at
        return True, True, False

    def _apply_autoscale(
        self,
        *,
        request: PlanResolutionRequest,
        now: float,
    ) -> tuple[PlanResolutionRequest, bool]:
        enabled = (
            request.autoscale_min_workers is not None
            or request.autoscale_max_workers is not None
            or request.autoscale_queue_depth is not None
            or request.autoscale_observed_latency_ms is not None
            or request.autoscale_observed_throughput_rps is not None
        )
        if not enabled:
            return request, False

        min_workers = request.autoscale_min_workers
        if min_workers is None:
            min_workers = 1
        max_workers = request.autoscale_max_workers
        if max_workers is None:
            max_workers = max(min_workers, request.num_partitions)
        min_workers = int(min_workers)
        max_workers = int(max_workers)
        if min_workers <= 0:
            raise ValueError("autoscale_min_workers must be > 0.")
        if max_workers < min_workers:
            raise ValueError("autoscale_max_workers must be >= autoscale_min_workers.")

        queue_target_size = request.autoscale_target_queue_per_worker
        if queue_target_size is None:
            queue_target_size = 32
        queue_target_size = int(queue_target_size)
        if queue_target_size <= 0:
            raise ValueError("autoscale_target_queue_per_worker must be > 0.")

        cooldown_ms = request.autoscale_cooldown_ms
        if cooldown_ms is None:
            cooldown_ms = 5000
        cooldown_ms = int(cooldown_ms)
        if cooldown_ms < 0:
            raise ValueError("autoscale_cooldown_ms must be >= 0.")

        stage_key = request.stage_id
        with self._metrics_lock:
            state = self._autoscale_states.get(stage_key)
            if state is None:
                initial = _clamp_int(request.num_partitions, min_workers, max_workers)
                state = _AutoscaleStageState(
                    desired_workers=initial,
                    cooldown_ms=cooldown_ms,
                )
                self._autoscale_states[stage_key] = state
            else:
                state.cooldown_ms = cooldown_ms

            current = _clamp_int(state.desired_workers, min_workers, max_workers)
            candidate = current
            reasons: list[str] = []

            if request.autoscale_queue_depth is not None:
                queue_depth = int(request.autoscale_queue_depth)
                if queue_depth < 0:
                    queue_depth = 0
                queue_target = int(math.ceil(float(queue_depth) / float(queue_target_size)))
                queue_target = _clamp_int(queue_target, min_workers, max_workers)
                if queue_target != candidate:
                    candidate = queue_target
                reasons.append("queue_depth")

            observed_latency = request.autoscale_observed_latency_ms
            latency_budget = request.autoscale_latency_budget_ms
            if (
                observed_latency is not None
                and latency_budget is not None
                and float(observed_latency) > float(latency_budget)
                and candidate < max_workers
            ):
                candidate = min(candidate + 1, max_workers)
                reasons.append("latency_pressure")

            observed_tput = request.autoscale_observed_throughput_rps
            target_tput = request.autoscale_target_throughput_rps
            if observed_tput is not None and target_tput is not None:
                observed_tput = float(observed_tput)
                target_tput = float(target_tput)
                queue_depth = int(request.autoscale_queue_depth or 0)
                if observed_tput < (target_tput * 0.9) and candidate < max_workers:
                    candidate = min(candidate + 1, max_workers)
                    reasons.append("throughput_deficit")
                elif (
                    observed_tput > (target_tput * 1.3)
                    and queue_depth == 0
                    and candidate > min_workers
                ):
                    candidate = max(candidate - 1, min_workers)
                    reasons.append("throughput_surplus")

            direction = "steady"
            if candidate > current:
                direction = "scale_up"
            elif candidate < current:
                direction = "scale_down"

            if direction != "steady" and cooldown_ms > 0:
                cooldown_s = float(cooldown_ms) / 1000.0
                if now < (state.last_scale_ts + cooldown_s):
                    candidate = current
                    direction = "steady"
                    reasons = ["cooldown_suppressed", *reasons]
                    state.cooldown_suppressed += 1
                    self._autoscale_cooldown_suppressed_count += 1

            if direction == "scale_up":
                self._autoscale_scale_up_count += 1
                state.last_scale_ts = now
            elif direction == "scale_down":
                self._autoscale_scale_down_count += 1
                state.last_scale_ts = now

            state.desired_workers = candidate
            state.last_direction = direction
            state.last_reason = ",".join(reasons) if reasons else "steady"
            state.last_decision_ts = now
            state.decisions += 1
            self._autoscale_decision_count += 1

            return replace(request, num_partitions=candidate), candidate != request.num_partitions

    def _autoscale_snapshot_locked(self) -> dict[str, object]:
        stage_snapshot: dict[str, dict[str, object]] = {}
        for stage_id, state in self._autoscale_states.items():
            stage_snapshot[stage_id] = {
                "desired_workers": state.desired_workers,
                "last_direction": state.last_direction,
                "last_reason": state.last_reason,
                "last_decision_ts": state.last_decision_ts,
                "cooldown_ms": state.cooldown_ms,
                "decisions": state.decisions,
                "cooldown_suppressed": state.cooldown_suppressed,
            }
        return {
            "decision_count": self._autoscale_decision_count,
            "scale_up_count": self._autoscale_scale_up_count,
            "scale_down_count": self._autoscale_scale_down_count,
            "cooldown_suppressed_count": self._autoscale_cooldown_suppressed_count,
            "active_stages": len(stage_snapshot),
            "stages": stage_snapshot,
        }

    def _count_locality_hits(self, plan: SchedulePlan) -> dict[str, int]:
        tier_by_address = self._tier_by_address()
        result: dict[str, int] = {}
        for address in plan.partitions.values():
            tier = tier_by_address.get(address, "unknown")
            result[tier] = result.get(tier, 0) + 1
        return result

    def _tier_by_address(self) -> dict[str, str]:
        mapping: dict[str, str] = {}
        for node in self._cluster_view.list_nodes():
            mapping.setdefault(node.address, self._cluster_view.locality_tier(node.node_id))
        return mapping

    def _record_metrics(
        self,
        *,
        latency_ms: float,
        staleness_ms: float,
        resource_fallback_used: bool,
        locality_hits: dict[str, int],
    ) -> None:
        safe_latency = max(0.0, latency_ms)
        safe_staleness = max(0.0, staleness_ms)
        with self._metrics_lock:
            self._plan_resolution_count += 1
            if resource_fallback_used:
                self._resource_fallback_count += 1

            self._staleness_count += 1
            self._staleness_ms_latest = safe_staleness
            self._staleness_ms_sum += safe_staleness
            if safe_staleness > self._staleness_ms_max:
                self._staleness_ms_max = safe_staleness

            self._latency_count += 1
            self._latency_ms_sum += safe_latency
            self._latency_ms_recent.append(safe_latency)
            if self._latency_ms_min is None or safe_latency < self._latency_ms_min:
                self._latency_ms_min = safe_latency
            if safe_latency > self._latency_ms_max:
                self._latency_ms_max = safe_latency

            for tier, count in locality_hits.items():
                if tier not in self._locality_hits:
                    self._locality_hits[tier] = 0
                self._locality_hits[tier] += count
                self._locality_total_hits += count


def _promote_node(nodes: list[ClusterNode], node_id: str) -> list[ClusterNode]:
    for idx, node in enumerate(nodes):
        if node.node_id == node_id:
            return [node, *nodes[:idx], *nodes[idx + 1 :]]
    return nodes


def _quantile(values: deque[float], q: float) -> float | None:
    if not values:
        return None
    if q <= 0.0:
        return min(values)
    if q >= 1.0:
        return max(values)
    samples = sorted(values)
    if len(samples) == 1:
        return float(samples[0])
    position = (len(samples) - 1) * float(q)
    lower = int(math.floor(position))
    upper = int(math.ceil(position))
    lower_value = float(samples[lower])
    upper_value = float(samples[upper])
    if lower == upper:
        return lower_value
    weight = position - float(lower)
    return lower_value + ((upper_value - lower_value) * weight)


def _dedupe_preserve_order(addresses: Iterable[str]) -> tuple[str, ...]:
    seen: set[str] = set()
    ordered: list[str] = []
    for address in addresses:
        if address in seen:
            continue
        seen.add(address)
        ordered.append(address)
    return tuple(ordered)


def _resource_match(
    summary: dict[str, object] | None,
    requirements: dict[str, int | float],
) -> bool:
    data = summary or {}
    for key, need in requirements.items():
        have = data.get(key)
        if not isinstance(have, (int, float)):
            return False
        if have < need:
            return False
    return True


def _owner_key(request: PlanResolutionRequest) -> str | None:
    if not request.lease_owner_kind or not request.lease_owner_id:
        return None
    return f"{request.lease_owner_kind}:{request.lease_owner_id}"


def _spillover_policy_enabled(request: PlanResolutionRequest) -> bool:
    return (
        request.spillover_max_remote_nodes is not None
        or request.spillover_ttl_ms is not None
        or not request.spillover_allow_remote
        or tuple(request.spillover_stage1_tiers) != ("local", "zone")
        or tuple(request.spillover_stage2_tiers) != ("region",)
        or request.spillover_claim_id is not None
    )


def _spillover_claim_key(request: PlanResolutionRequest) -> str:
    if request.spillover_claim_id:
        return request.spillover_claim_id
    owner_key = _owner_key(request)
    if owner_key:
        return f"{request.stage_id}:{owner_key}"
    return f"{request.stage_id}:{request.request_id}"


def _clamp_int(value: int, lower: int, upper: int) -> int:
    return max(lower, min(value, upper))


__all__ = [
    "DEFAULT_LOCALITY_ORDER",
    "PlanResolutionRequest",
    "PlanResolution",
    "ClusterPlanResolver",
]
