prole/knoe/core/cluster_storage_browser.py

from __future__ import annotations

import json
import re
import subprocess
import time
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from typing import Protocol


PD_SSD_QUOTA_METRIC = "SSD_TOTAL_GB"
PD_BALANCED_QUOTA_METRIC = "DISKS_TOTAL_GB"

STORAGE_CLASS_PREMIUM = "premium-rwo"
STORAGE_CLASS_STANDARD = "standard-rwo"

_REGION_TOKEN_RE = re.compile(r"^[a-z]+-[a-z]+\d+$")


@dataclass(slots=True, frozen=True)
class RegionCatalogEntry:
    region: str
    timezone_group: str
    country: str


@dataclass(slots=True, frozen=True)
class QuotaMetricHeadroom:
    metric: str
    limit_gb: float | None
    usage_gb: float | None

    @property
    def headroom_gb(self) -> float | None:
        if self.limit_gb is None or self.usage_gb is None:
            return None
        return self.limit_gb - self.usage_gb


@dataclass(slots=True, frozen=True)
class RegionQuotaSnapshot:
    region: str
    premium: QuotaMetricHeadroom
    standard: QuotaMetricHeadroom


@dataclass(slots=True, frozen=True)
class RegionLatencyResult:
    region: str
    latency_ms: float | None
    ok: bool
    error: str = ""


@dataclass(slots=True, frozen=True)
class ClusterStorageRequest:
    pgdata_gb: float
    wal_gb: float

    @property
    def total_gb(self) -> float:
        return max(0.0, self.pgdata_gb) + max(0.0, self.wal_gb)


@dataclass(slots=True, frozen=True)
class RegionFeasibility:
    region: str
    timezone_group: str
    premium_headroom_gb: float | None
    standard_headroom_gb: float | None
    can_premium: bool
    can_standard: bool
    default_pgdata_class: str
    default_wal_class: str
    feasibility_rank: int
    notes: tuple[str, ...]


@dataclass(slots=True, frozen=True)
class RegionCandidate:
    feasibility: RegionFeasibility
    latency_ms: float | None


@dataclass(slots=True, frozen=True)
class ClusterStorageBrowserResult:
    generated_at: str
    project_id: str
    timezone_group: str
    request: ClusterStorageRequest
    candidates: tuple[RegionCandidate, ...]
    quota_link: str
    quota_console_link: str

    def to_dict(self) -> dict:
        return {
            "generatedAt": self.generated_at,
            "projectId": self.project_id,
            "timezoneGroup": self.timezone_group,
            "request": asdict(self.request),
            "candidates": [
                {
                    "region": c.feasibility.region,
                    "timezoneGroup": c.feasibility.timezone_group,
                    "feasibilityRank": c.feasibility.feasibility_rank,
                    "premiumHeadroomGb": c.feasibility.premium_headroom_gb,
                    "standardHeadroomGb": c.feasibility.standard_headroom_gb,
                    "canPremium": c.feasibility.can_premium,
                    "canStandard": c.feasibility.can_standard,
                    "defaultPgdataClass": c.feasibility.default_pgdata_class,
                    "defaultWalClass": c.feasibility.default_wal_class,
                    "latencyMs": c.latency_ms,
                    "notes": list(c.feasibility.notes),
                }
                for c in self.candidates
            ],
            "quotaLink": self.quota_link,
            "quotaConsoleLink": self.quota_console_link,
        }

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":"), indent=2)


class RegionQuotaProvider(Protocol):
    def fetch(self, *, project_id: str, regions: list[str]) -> dict[str, RegionQuotaSnapshot]: ...


class RegionLatencyProvider(Protocol):
    def probe(self, regions: list[str]) -> dict[str, RegionLatencyResult]: ...


def default_region_catalog() -> tuple[RegionCatalogEntry, ...]:
    return (
        RegionCatalogEntry("us-central1", "americas", "us"),
        RegionCatalogEntry("us-east1", "americas", "us"),
        RegionCatalogEntry("us-east4", "americas", "us"),
        RegionCatalogEntry("us-west1", "americas", "us"),
        RegionCatalogEntry("us-west2", "americas", "us"),
        RegionCatalogEntry("us-west3", "americas", "us"),
        RegionCatalogEntry("us-west4", "americas", "us"),
        RegionCatalogEntry("northamerica-northeast1", "americas", "ca"),
        RegionCatalogEntry("southamerica-east1", "americas", "br"),
        RegionCatalogEntry("europe-west1", "emea", "be"),
        RegionCatalogEntry("europe-west2", "emea", "uk"),
        RegionCatalogEntry("europe-west3", "emea", "de"),
        RegionCatalogEntry("europe-west4", "emea", "nl"),
        RegionCatalogEntry("europe-west6", "emea", "ch"),
        RegionCatalogEntry("europe-west8", "emea", "it"),
        RegionCatalogEntry("europe-west9", "emea", "fr"),
        RegionCatalogEntry("europe-west10", "emea", "de"),
        RegionCatalogEntry("europe-central2", "emea", "pl"),
        RegionCatalogEntry("asia-east1", "apac", "tw"),
        RegionCatalogEntry("asia-east2", "apac", "hk"),
        RegionCatalogEntry("asia-northeast1", "apac", "jp"),
        RegionCatalogEntry("asia-northeast2", "apac", "jp"),
        RegionCatalogEntry("asia-northeast3", "apac", "kr"),
        RegionCatalogEntry("asia-south1", "apac", "in"),
        RegionCatalogEntry("asia-south2", "apac", "in"),
        RegionCatalogEntry("asia-southeast1", "apac", "sg"),
        RegionCatalogEntry("asia-southeast2", "apac", "id"),
        RegionCatalogEntry("australia-southeast1", "apac", "au"),
        RegionCatalogEntry("australia-southeast2", "apac", "au"),
    )


def detect_timezone_group() -> str:
    now_local = datetime.now().astimezone()
    offset_hours = float(now_local.utcoffset().total_seconds()) / 3600.0 if now_local.utcoffset() else 0.0
    if -10.0 <= offset_hours <= -2.0:
        return "americas"
    if -1.0 <= offset_hours <= 4.0:
        return "emea"
    return "apac"


def filter_regions_by_timezone(
    regions: list[str],
    timezone_group: str,
    catalog: tuple[RegionCatalogEntry, ...] | None = None,
) -> list[RegionCatalogEntry]:
    entries = catalog or default_region_catalog()
    allowed = {r for r in regions if _REGION_TOKEN_RE.match(r)}
    return [entry for entry in entries if entry.region in allowed and entry.timezone_group == timezone_group]


def _bool_feasible(headroom_gb: float | None, requested_gb: float) -> bool:
    if headroom_gb is None:
        return False
    return headroom_gb >= max(0.0, requested_gb)


def _choose_default_class(premium_ok: bool, standard_ok: bool) -> str:
    if premium_ok:
        return STORAGE_CLASS_PREMIUM
    if standard_ok:
        return STORAGE_CLASS_STANDARD
    return STORAGE_CLASS_PREMIUM


def _feasibility_rank(can_premium: bool, can_standard: bool) -> int:
    if can_premium:
        return 0
    if can_standard:
        return 1
    return 2


def evaluate_region_feasibility(
    region: str,
    timezone_group: str,
    quota: RegionQuotaSnapshot | None,
    requested: ClusterStorageRequest,
) -> RegionFeasibility:
    premium_headroom = quota.premium.headroom_gb if quota is not None else None
    standard_headroom = quota.standard.headroom_gb if quota is not None else None

    can_premium = _bool_feasible(premium_headroom, requested.total_gb)
    can_standard = _bool_feasible(standard_headroom, requested.total_gb)

    notes: list[str] = []
    if quota is None:
        notes.append("quota-unavailable")
    if not can_premium and not can_standard:
        notes.append("insufficient-quota")

    default_class = _choose_default_class(can_premium, can_standard)
    return RegionFeasibility(
        region=region,
        timezone_group=timezone_group,
        premium_headroom_gb=premium_headroom,
        standard_headroom_gb=standard_headroom,
        can_premium=can_premium,
        can_standard=can_standard,
        default_pgdata_class=default_class,
        default_wal_class=default_class,
        feasibility_rank=_feasibility_rank(can_premium, can_standard),
        notes=tuple(notes),
    )


def _quota_links(project_id: str) -> tuple[str, str]:
    project = project_id.strip()
    if not project:
        return (
            "https://cloud.google.com/compute/quotas",
            "https://console.cloud.google.com/iam-admin/quotas",
        )
    return (
        "https://cloud.google.com/compute/quotas",
        f"https://console.cloud.google.com/iam-admin/quotas?project={project}",
    )


def _sort_candidates(candidates: list[RegionCandidate]) -> list[RegionCandidate]:
    def _key(candidate: RegionCandidate):
        latency = candidate.latency_ms if candidate.latency_ms is not None else float("inf")
        return (candidate.feasibility.feasibility_rank, latency, candidate.feasibility.region)

    return sorted(candidates, key=_key)


def build_cluster_storage_browser_result(
    *,
    project_id: str,
    requested: ClusterStorageRequest,
    available_regions: list[str],
    timezone_group: str | None = None,
    catalog: tuple[RegionCatalogEntry, ...] | None = None,
    quota_provider: RegionQuotaProvider,
    latency_provider: RegionLatencyProvider,
) -> ClusterStorageBrowserResult:
    tz_group = (timezone_group or detect_timezone_group()).strip().lower() or detect_timezone_group()
    candidates = filter_regions_by_timezone(available_regions, tz_group, catalog)
    candidate_regions = [entry.region for entry in candidates]

    quota_by_region = quota_provider.fetch(project_id=project_id, regions=candidate_regions)
    feasibility: list[RegionFeasibility] = [
        evaluate_region_feasibility(
            region=entry.region,
            timezone_group=entry.timezone_group,
            quota=quota_by_region.get(entry.region),
            requested=requested,
        )
        for entry in candidates
    ]

    probe_regions = [f.region for f in feasibility if f.can_premium or f.can_standard]
    latency_by_region = latency_provider.probe(probe_regions)

    combined: list[RegionCandidate] = []
    for f in feasibility:
        latency = latency_by_region.get(f.region)
        combined.append(
            RegionCandidate(
                feasibility=f,
                latency_ms=latency.latency_ms if latency and latency.ok else None,
            )
        )

    link, console_link = _quota_links(project_id)
    return ClusterStorageBrowserResult(
        generated_at=datetime.now(timezone.utc).isoformat(),
        project_id=project_id,
        timezone_group=tz_group,
        request=requested,
        candidates=tuple(_sort_candidates(combined)),
        quota_link=link,
        quota_console_link=console_link,
    )


def _safe_float(value: str | int | float | None) -> float | None:
    if value is None:
        return None
    try:
        return float(value)
    except Exception:
        return None


def _parse_limit_usage(raw: str) -> tuple[float | None, float | None]:
    # gcloud output format from --format=value(limit,usage)
    parts = [p.strip() for p in raw.split() if p.strip()]
    if len(parts) >= 2:
        return _safe_float(parts[0]), _safe_float(parts[1])
    if len(parts) == 1:
        return _safe_float(parts[0]), None
    return None, None


class GcloudRegionQuotaProvider:
    def __init__(self, timeout_sec: int = 20):
        self.timeout_sec = timeout_sec

    def _metric_headroom(self, project_id: str, region: str, metric: str) -> QuotaMetricHeadroom:
        cmd = [
            "gcloud",
            "compute",
            "regions",
            "describe",
            region,
            "--project",
            project_id,
            "--format=value(quotas[metric=%s].limit,quotas[metric=%s].usage)" % (metric, metric),
        ]
        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=self.timeout_sec,
            )
            if result.returncode != 0:
                return QuotaMetricHeadroom(metric=metric, limit_gb=None, usage_gb=None)
            limit, usage = _parse_limit_usage(result.stdout.strip())
            return QuotaMetricHeadroom(metric=metric, limit_gb=limit, usage_gb=usage)
        except Exception:
            return QuotaMetricHeadroom(metric=metric, limit_gb=None, usage_gb=None)

    def fetch(self, *, project_id: str, regions: list[str]) -> dict[str, RegionQuotaSnapshot]:
        project = project_id.strip()
        if not project:
            return {}
        out: dict[str, RegionQuotaSnapshot] = {}
        for region in regions:
            premium = self._metric_headroom(project, region, PD_SSD_QUOTA_METRIC)
            standard = self._metric_headroom(project, region, PD_BALANCED_QUOTA_METRIC)
            out[region] = RegionQuotaSnapshot(region=region, premium=premium, standard=standard)
        return out


class PingLatencyProvider:
    def __init__(self, timeout_sec: int = 2, count: int = 2):
        self.timeout_sec = timeout_sec
        self.count = count

    def _probe_one(self, region: str) -> RegionLatencyResult:
        # Best-effort: use region DNS style endpoint as latency anchor.
        host = f"{region}.gcping.com"
        cmd = ["ping", "-c", str(self.count), "-t", str(self.timeout_sec), host]
        started = time.monotonic()
        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=max(2, self.timeout_sec + 2),
            )
            elapsed = (time.monotonic() - started) * 1000.0
            if result.returncode == 0:
                return RegionLatencyResult(region=region, latency_ms=round(elapsed, 2), ok=True)
            return RegionLatencyResult(region=region, latency_ms=None, ok=False, error="ping-failed")
        except Exception as exc:
            return RegionLatencyResult(region=region, latency_ms=None, ok=False, error=str(exc))

    def probe(self, regions: list[str]) -> dict[str, RegionLatencyResult]:
        return {region: self._probe_one(region) for region in regions}


class StaticQuotaProvider:
    def __init__(self, values: dict[str, RegionQuotaSnapshot]):
        self.values = dict(values)

    def fetch(self, *, project_id: str, regions: list[str]) -> dict[str, RegionQuotaSnapshot]:
        _ = project_id
        return {r: self.values[r] for r in regions if r in self.values}


class StaticLatencyProvider:
    def __init__(self, values: dict[str, RegionLatencyResult]):
        self.values = dict(values)

    def probe(self, regions: list[str]) -> dict[str, RegionLatencyResult]:
        return {r: self.values[r] for r in regions if r in self.values}