from __future__ import annotations

import argparse
import json
import math
import os
import sys
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from datetime import datetime
from datetime import timezone
from typing import Any


DEFAULT_BASE_URL = "https://clinic.pingcap.com"
DEFAULT_TIMEOUT_S = 60
TARGET_ROLES = ("tikv", "tiflash")


class MetricsQueryError(RuntimeError):
    pass


@dataclass(frozen=True)
class StoreRoleMapping:
    store: str
    pod: str
    component: str
    address: str


class ClinicMetricsClient:
    def __init__(self, *, api_key: str, base_url: str, timeout_s: int = DEFAULT_TIMEOUT_S) -> None:
        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._timeout_s = timeout_s

    def query_instant(
        self,
        *,
        cluster_id: str,
        promql: str,
        query_time: int,
    ) -> dict[str, Any]:
        return self._get(
            f"/data-proxy/metrics/{cluster_id}/api/v1/query",
            params={
                "query": promql,
                "time": str(query_time),
            },
        )

    def _get(self, path: str, *, params: dict[str, str]) -> dict[str, Any]:
        url = f"{self._base_url}{path}?{urllib.parse.urlencode(params)}"
        request = urllib.request.Request(
            url,
            headers={"Authorization": f"Bearer {self._api_key}"},
            method="GET",
        )
        try:
            with urllib.request.urlopen(request, timeout=self._timeout_s) as response:
                payload = response.read().decode("utf-8")
        except urllib.error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")
            raise MetricsQueryError(f"HTTP {exc.code} for {path}: {body}") from exc
        except urllib.error.URLError as exc:
            raise MetricsQueryError(f"request failed for {path}: {exc}") from exc
        try:
            data = json.loads(payload)
        except json.JSONDecodeError as exc:
            raise MetricsQueryError(f"invalid JSON response for {path}: {exc}") from exc
        status = str(data.get("status") or "")
        if status and status != "success":
            raise MetricsQueryError(f"unexpected response status for {path}: {status}")
        return data


def _env_api_key() -> str:
    for key in ("SKILL_CLINIC_API_KEY", "CLINIC_API_KEY", "INVESTIBOOK_CLINIC_API_KEY"):
        value = os.environ.get(key, "").strip()
        if value:
            return value
    raise MetricsQueryError(
        "missing Clinic API key in SKILL_CLINIC_API_KEY/CLINIC_API_KEY/INVESTIBOOK_CLINIC_API_KEY"
    )


def _isoformat_utc(ts: int) -> str:
    return datetime.fromtimestamp(ts, tz=timezone.utc).isoformat().replace("+00:00", "Z")


def _parse_iso8601(value: str) -> int:
    normalized = value.strip()
    if normalized.endswith("Z"):
        normalized = normalized[:-1] + "+00:00"
    return int(datetime.fromisoformat(normalized).timestamp())


def _safe_float(value: Any) -> float | None:
    try:
        numeric = float(value)
    except (TypeError, ValueError):
        return None
    if not math.isfinite(numeric):
        return None
    return numeric


def _result_series(payload: dict[str, Any]) -> list[dict[str, Any]]:
    data = payload.get("data")
    if not isinstance(data, dict):
        return []
    result = data.get("result")
    return [item for item in result if isinstance(item, dict)] if isinstance(result, list) else []


def _namespace_for_cluster(cluster_id: str) -> str:
    return f"tidb{cluster_id}"


def _active_pods_query(cluster_id: str) -> str:
    namespace = _namespace_for_cluster(cluster_id)
    return (
        f'count(kube_pod_info{{namespace="{namespace}",pod=~"db-.*(tikv|tiflash).*"}}) by (pod)'
    )


def _pd_store_address_query() -> str:
    return "count(pd_scheduler_store_status) by (address, store)"


def _align_snapshot_time(*, now_ts: int, step_seconds: int, lag_seconds: int) -> int:
    effective_now = now_ts - lag_seconds
    if effective_now <= 0:
        raise MetricsQueryError("effective current time must be positive")
    snapshot_ts = (effective_now // step_seconds) * step_seconds
    if snapshot_ts <= 0:
        raise MetricsQueryError("resolved snapshot time must be positive")
    return snapshot_ts


def _storage_pod_sort_key(pod: str) -> tuple[int, str]:
    try:
        suffix = int(pod.rsplit("-", 1)[1])
    except (IndexError, ValueError):
        return (sys.maxsize, pod)
    return (suffix, pod)


def _active_storage_pods(payload: dict[str, Any]) -> set[str]:
    pods: set[str] = set()
    for item in _result_series(payload):
        metric = item.get("metric")
        if not isinstance(metric, dict):
            continue
        pod = str(metric.get("pod") or "").strip()
        if pod:
            pods.add(pod)
    if not pods:
        raise MetricsQueryError("resolved empty active storage pod inventory")
    return pods


def _component_for_pod(pod: str) -> str:
    if pod.startswith("db-tikv-"):
        return "tikv"
    if pod.startswith("db-tiflash-"):
        return "tiflash"
    raise MetricsQueryError(f"unrecognized storage pod role for pod '{pod}'")


def _active_store_role_mappings(
    payload: dict[str, Any],
    *,
    active_pods: set[str],
) -> list[StoreRoleMapping]:
    pod_to_mapping: dict[str, StoreRoleMapping] = {}
    store_to_pod: dict[str, str] = {}
    for item in _result_series(payload):
        metric = item.get("metric")
        if not isinstance(metric, dict):
            continue
        store = str(metric.get("store") or "").strip()
        address = str(metric.get("address") or "").strip()
        if not store or not address:
            continue
        pod = address.split(".", 1)[0].strip()
        if pod not in active_pods:
            continue
        mapping = StoreRoleMapping(
            store=store,
            pod=pod,
            component=_component_for_pod(pod),
            address=address,
        )
        previous = pod_to_mapping.get(pod)
        if previous is not None and previous.store != mapping.store:
            raise MetricsQueryError(
                f"active pod '{pod}' maps to multiple stores: '{previous.store}' and '{mapping.store}'"
            )
        other_pod = store_to_pod.get(store)
        if other_pod is not None and other_pod != pod:
            raise MetricsQueryError(
                f"store '{store}' maps to multiple active pods: '{other_pod}' and '{pod}'"
            )
        pod_to_mapping[pod] = mapping
        store_to_pod[store] = pod

    missing_pods = sorted(active_pods - set(pod_to_mapping))
    if missing_pods:
        raise MetricsQueryError(
            f"missing PD store mapping for active storage pods: {', '.join(missing_pods[:10])}"
        )
    return sorted(
        pod_to_mapping.values(),
        key=lambda item: (item.component, _storage_pod_sort_key(item.pod)),
    )


def _spread_indices(length: int, count: int) -> list[int]:
    if length <= 0 or count <= 0:
        return []
    if count >= length:
        return list(range(length))
    raw = [round(idx * (length - 1) / (count - 1)) if count > 1 else 0 for idx in range(count)]
    indices: list[int] = []
    seen: set[int] = set()
    for value in raw:
        bounded = min(max(int(value), 0), length - 1)
        if bounded in seen:
            continue
        seen.add(bounded)
        indices.append(bounded)
    cursor = 0
    while len(indices) < count and cursor < length:
        if cursor not in seen:
            seen.add(cursor)
            indices.append(cursor)
        cursor += 1
    return sorted(indices)


def _select_store_samples(
    mappings: list[StoreRoleMapping],
    *,
    stores_per_role: int,
) -> list[StoreRoleMapping]:
    selected_by_role: dict[str, list[StoreRoleMapping]] = {}
    for role in TARGET_ROLES:
        role_items = [item for item in mappings if item.component == role]
        if len(role_items) < stores_per_role:
            raise MetricsQueryError(
                f"expected at least {stores_per_role} active {role} stores, got {len(role_items)}"
            )
        indices = _spread_indices(len(role_items), stores_per_role)
        selected_by_role[role] = [role_items[idx] for idx in indices]

    selected: list[StoreRoleMapping] = []
    for idx in range(stores_per_role):
        for role in TARGET_ROLES:
            selected.append(selected_by_role[role][idx])
    return selected


def _build_case_input(args: argparse.Namespace) -> dict[str, Any]:
    client = ClinicMetricsClient(
        api_key=_env_api_key(),
        base_url=os.environ.get("CLINIC_ENDPOINT", DEFAULT_BASE_URL).strip() or DEFAULT_BASE_URL,
    )
    snapshot_ts = _align_snapshot_time(
        now_ts=int(datetime.now(tz=timezone.utc).timestamp()),
        step_seconds=args.step_seconds,
        lag_seconds=args.snapshot_lag_seconds,
    )
    active_payload = client.query_instant(
        cluster_id=args.cluster_id,
        promql=_active_pods_query(args.cluster_id),
        query_time=snapshot_ts,
    )
    active_pods = _active_storage_pods(active_payload)
    pd_payload = client.query_instant(
        cluster_id=args.cluster_id,
        promql=_pd_store_address_query(),
        query_time=snapshot_ts,
    )
    mappings = _active_store_role_mappings(pd_payload, active_pods=active_pods)
    selected = _select_store_samples(mappings, stores_per_role=args.stores_per_role)
    store_ids = [item.store for item in selected]
    return {
        "id": f"snapshot-{datetime.fromtimestamp(snapshot_ts, tz=timezone.utc).strftime('%Y%m%dT%H%M%SZ')}",
        "cluster_id": args.cluster_id,
        "start": _isoformat_utc(snapshot_ts - args.step_seconds),
        "end": _isoformat_utc(snapshot_ts),
        "snapshot_time": _isoformat_utc(snapshot_ts),
        "step_seconds": args.step_seconds,
        "stores_per_role": args.stores_per_role,
        "store_ids": store_ids,
        "store_ids_csv": ",".join(store_ids),
        "store_ids_json": json.dumps(store_ids, ensure_ascii=False),
        "selected_stores": [
            {
                "store": item.store,
                "component": item.component,
                "pod": item.pod,
                "address": item.address,
            }
            for item in selected
        ],
        "inventory_summary": {
            "tikv": sum(1 for item in mappings if item.component == "tikv"),
            "tiflash": sum(1 for item in mappings if item.component == "tiflash"),
        },
        "queries": {
            "active_pods_query": _active_pods_query(args.cluster_id),
            "pd_store_address_query": _pd_store_address_query(),
        },
    }


def _build_reference(args: argparse.Namespace) -> dict[str, Any]:
    client = ClinicMetricsClient(
        api_key=_env_api_key(),
        base_url=os.environ.get("CLINIC_ENDPOINT", DEFAULT_BASE_URL).strip() or DEFAULT_BASE_URL,
    )
    snapshot_ts = _parse_iso8601(args.snapshot_time)
    store_ids = [item.strip() for item in str(args.stores_csv or "").split(",") if item.strip()]
    if not store_ids:
        raise MetricsQueryError("resolved empty selected store list")
    active_payload = client.query_instant(
        cluster_id=args.cluster_id,
        promql=_active_pods_query(args.cluster_id),
        query_time=snapshot_ts,
    )
    active_pods = _active_storage_pods(active_payload)
    pd_payload = client.query_instant(
        cluster_id=args.cluster_id,
        promql=_pd_store_address_query(),
        query_time=snapshot_ts,
    )
    mappings = _active_store_role_mappings(pd_payload, active_pods=active_pods)
    by_store = {item.store: item for item in mappings}
    stores: list[dict[str, str]] = []
    for store in store_ids:
        mapping = by_store.get(store)
        if mapping is None:
            raise MetricsQueryError(f"selected store '{store}' was not found in the active store inventory")
        stores.append(
            {
                "store": mapping.store,
                "component": mapping.component,
            }
        )
    return {
        "cluster_id": args.cluster_id,
        "evidence_sufficient": True,
        "stores": stores,
        "snapshot_time": _isoformat_utc(snapshot_ts),
    }


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Runtime hook for store role classification case")
    subparsers = parser.add_subparsers(dest="command", required=True)

    input_parser = subparsers.add_parser("prepare-case-input")
    input_parser.add_argument("--cluster-id", required=True)
    input_parser.add_argument("--snapshot-lag-seconds", type=int, default=600)
    input_parser.add_argument("--step-seconds", type=int, default=300)
    input_parser.add_argument("--stores-per-role", type=int, default=3)

    reference_parser = subparsers.add_parser("build-reference")
    reference_parser.add_argument("--cluster-id", required=True)
    reference_parser.add_argument("--snapshot-time", required=True)
    reference_parser.add_argument("--stores-csv", required=True)

    return parser


def main(argv: list[str]) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)
    if args.command == "prepare-case-input":
        payload = _build_case_input(args)
    elif args.command == "build-reference":
        payload = _build_reference(args)
    else:
        raise MetricsQueryError(f"unsupported command: {args.command}")
    sys.stdout.write(json.dumps(payload, ensure_ascii=False))
    sys.stdout.write("\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))