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 UTC
from datetime import datetime
from datetime import timedelta
from typing import Any
from zoneinfo import ZoneInfo


DEFAULT_BASE_URL = "https://clinic.pingcap.com"
DEFAULT_TIMEOUT_S = 60
BEIJING_TZ = ZoneInfo("Asia/Shanghai")


@dataclass(frozen=True)
class TiKVNodeInventory:
    node: str
    tikv_node: str


class MetricsQueryError(RuntimeError):
    pass


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_range(
        self,
        *,
        cluster_id: str,
        promql: str,
        start: int,
        end: int,
        step_seconds: int,
    ) -> dict[str, Any]:
        return self._get(
            f"/data-proxy/metrics/{cluster_id}/api/v1/query_range",
            params={
                "query": promql,
                "start": str(start),
                "end": str(end),
                "step": str(step_seconds),
            },
        )

    def query_series(
        self,
        *,
        cluster_id: str,
        match: str,
        start: int,
        end: int,
    ) -> dict[str, Any]:
        return self._get(
            f"/data-proxy/metrics/{cluster_id}/api/v1/series",
            params={
                "match[]": [match],
                "start": str(start),
                "end": str(end),
            },
        )

    def _get(self, path: str, *, params: dict[str, Any]) -> dict[str, Any]:
        url = f"{self._base_url}{path}?{urllib.parse.urlencode(params, doseq=True)}"
        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 _parse_iso8601(value: str) -> int:
    normalized = value.strip()
    if normalized.endswith("Z"):
        normalized = normalized[:-1] + "+00:00"
    return int(datetime.fromisoformat(normalized).timestamp())


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


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


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


def _disk_read_query(instance_regex: str, *, rate_window_seconds: int) -> str:
    selector = f'node_disk_reads_completed_total{{instance=~"{instance_regex}",device!~"dm.*"}}'
    return f"sum by (instance) (rate({selector}[{rate_window_seconds}s]))"


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 _series_discovery_items(payload: dict[str, Any]) -> list[dict[str, Any]]:
    data = payload.get("data")
    return [item for item in data if isinstance(item, dict)] if isinstance(data, list) else []


def _inventory_pairs(payload: dict[str, Any]) -> list[TiKVNodeInventory]:
    resolved: dict[str, str] = {}
    for item in _series_discovery_items(payload):
        node = str(item.get("node") or "").strip()
        tikv_node = str(item.get("pod") or "").strip()
        if not node or not tikv_node:
            continue
        previous = resolved.get(node)
        if previous is None:
            resolved[node] = tikv_node
            continue
        if previous != tikv_node:
            raise MetricsQueryError(
                f"node '{node}' maps to multiple TiKV pods: '{previous}' and '{tikv_node}'"
            )
    return [
        TiKVNodeInventory(node=node, tikv_node=tikv_node)
        for node, tikv_node in sorted(resolved.items())
    ]


def _instance_regex(instances: list[str]) -> str:
    unique_instances = [instance for instance in sorted(set(instances)) if instance]
    if not unique_instances:
        raise MetricsQueryError("resolved empty TiKV instance inventory")
    return "^(?:" + "|".join(instance.replace(".", "[.]") for instance in unique_instances) + ")$"


def _series_max_value(item: dict[str, Any]) -> float | None:
    values = item.get("values")
    if not isinstance(values, list):
        return None
    observed: list[float] = []
    for row in values:
        if not isinstance(row, list) or len(row) < 2:
            continue
        parsed = _safe_float(row[1])
        if parsed is not None:
            observed.append(parsed)
    return max(observed) if observed else None


def _build_cluster_rows(
    client: ClinicMetricsClient,
    *,
    cluster_id: str,
    start_ts: int,
    end_ts: int,
    step_seconds: int,
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    inventory_payload = client.query_series(
        cluster_id=cluster_id,
        match=_inventory_match(cluster_id),
        start=start_ts,
        end=end_ts,
    )
    inventory = _inventory_pairs(inventory_payload)
    if not inventory:
        raise MetricsQueryError(f"failed to resolve TiKV node inventory for cluster {cluster_id}")

    node_to_tikv_node = {item.node: item.tikv_node for item in inventory}
    sample_start_ts = start_ts + 60
    if sample_start_ts > end_ts:
        raise MetricsQueryError("window is too short to compute max 1-minute read IOPS")
    range_payload = client.query_range(
        cluster_id=cluster_id,
        promql=_disk_read_query(_instance_regex(list(node_to_tikv_node)), rate_window_seconds=60),
        start=sample_start_ts,
        end=end_ts,
        step_seconds=step_seconds,
    )

    rows: list[dict[str, Any]] = []
    observed_instances: set[str] = set()
    for item in _result_series(range_payload):
        metric = item.get("metric")
        if not isinstance(metric, dict):
            continue
        instance = str(metric.get("instance") or "").strip()
        if not instance:
            continue
        tikv_node = node_to_tikv_node.get(instance)
        if not tikv_node:
            raise MetricsQueryError(
                f"metric instance '{instance}' was not found in resolved TiKV node inventory for cluster {cluster_id}"
            )
        max_value = _series_max_value(item)
        if max_value is None:
            continue
        observed_instances.add(instance)
        rows.append(
            {
                "cluster_id": cluster_id,
                "tikv_node": tikv_node,
                "max_1m_read_iops": round(max_value, 3),
            }
        )

    missing_instances = sorted(set(node_to_tikv_node) - observed_instances)
    if missing_instances:
        missing_tikv_nodes = [node_to_tikv_node[instance] for instance in missing_instances]
        raise MetricsQueryError(
            "missing read IOPS series for "
            f"{len(missing_instances)} TiKV nodes in cluster {cluster_id}: "
            + ", ".join(missing_tikv_nodes)
        )

    rows.sort(key=lambda item: (-float(item["max_1m_read_iops"]), str(item["tikv_node"])))
    return rows, {
        "inventory_match": _inventory_match(cluster_id),
        "inventory_count": len(inventory),
        "ranked_count": len(rows),
        "range_query": _disk_read_query(_instance_regex(list(node_to_tikv_node)), rate_window_seconds=60),
        "sample_start": _isoformat_utc(sample_start_ts),
        "sample_end": _isoformat_utc(end_ts),
    }


def _latest_completed_beijing_window(
    *,
    local_hour: int,
    local_minute: int,
    duration_minutes: int,
) -> tuple[int, int, str]:
    now_local = datetime.now(tz=BEIJING_TZ)
    candidate_start = now_local.replace(
        hour=local_hour,
        minute=local_minute,
        second=0,
        microsecond=0,
    )
    candidate_end = candidate_start + timedelta(minutes=duration_minutes)
    if now_local < candidate_end:
        candidate_start = candidate_start - timedelta(days=1)
        candidate_end = candidate_start + timedelta(minutes=duration_minutes)
    return (
        int(candidate_start.astimezone(UTC).timestamp()),
        int(candidate_end.astimezone(UTC).timestamp()),
        candidate_start.date().isoformat(),
    )


def _resolve_window(args: argparse.Namespace) -> dict[str, Any]:
    window_start, window_end, local_date = _latest_completed_beijing_window(
        local_hour=args.local_hour,
        local_minute=args.local_minute,
        duration_minutes=args.duration_minutes,
    )
    cluster_ids = sorted([args.cluster_id_a, args.cluster_id_b])
    return {
        "id": f"bjt-{local_date}-{args.local_hour:02d}{args.local_minute:02d}-{args.duration_minutes}m",
        "cluster_ids": cluster_ids,
        "start": _isoformat_utc(window_start),
        "end": _isoformat_utc(window_end),
        "local_timezone": "Asia/Shanghai",
        "local_date": local_date,
        "local_hour": args.local_hour,
        "local_minute": args.local_minute,
        "duration_minutes": args.duration_minutes,
        "step_seconds": args.step_seconds,
        "mapping_strategy": "node_filter",
    }


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,
    )
    start_ts = _parse_iso8601(args.start)
    end_ts = _parse_iso8601(args.end)
    if end_ts <= start_ts:
        raise MetricsQueryError("reference window end must be greater than start")

    cluster_ids = sorted([args.cluster_id_a, args.cluster_id_b])
    cluster_queries: dict[str, Any] = {}
    cluster_node_counts: dict[str, int] = {}
    grouped_rows: list[dict[str, Any]] = []
    for cluster_id in cluster_ids:
        cluster_rows, cluster_debug = _build_cluster_rows(
            client,
            cluster_id=cluster_id,
            start_ts=start_ts,
            end_ts=end_ts,
            step_seconds=args.step_seconds,
        )
        cluster_queries[cluster_id] = cluster_debug
        cluster_node_counts[cluster_id] = len(cluster_rows)
        grouped_rows.append(
            {
                "cluster_id": cluster_id,
                "nodes": [
                    [str(item["tikv_node"]), float(item["max_1m_read_iops"])]
                    for item in cluster_rows
                ],
            }
        )

    grouped_rows.sort(key=lambda item: str(item["cluster_id"]))
    return {
        "window": {
            "start": _isoformat_utc(start_ts),
            "end": _isoformat_utc(end_ts),
        },
        "mapping_strategy": "node_filter",
        "evidence_sufficient": True,
        "cluster_ids": cluster_ids,
        "cluster_node_counts": cluster_node_counts,
        "nodes_by_cluster": grouped_rows,
        "queries": cluster_queries,
    }


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Runtime hook for Candidate B full inventory case")
    subparsers = parser.add_subparsers(dest="command", required=True)

    resolve_parser = subparsers.add_parser("resolve-window")
    resolve_parser.add_argument("--cluster-id-a", required=True)
    resolve_parser.add_argument("--cluster-id-b", required=True)
    resolve_parser.add_argument("--local-hour", type=int, default=11)
    resolve_parser.add_argument("--local-minute", type=int, default=0)
    resolve_parser.add_argument("--duration-minutes", type=int, default=5)
    resolve_parser.add_argument("--step-seconds", type=int, default=60)

    reference_parser = subparsers.add_parser("build-reference")
    reference_parser.add_argument("--cluster-id-a", required=True)
    reference_parser.add_argument("--cluster-id-b", required=True)
    reference_parser.add_argument("--start", required=True)
    reference_parser.add_argument("--end", required=True)
    reference_parser.add_argument("--step-seconds", type=int, default=60)

    return parser


def main(argv: list[str]) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)
    if args.command == "resolve-window":
        payload = _resolve_window(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:]))