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

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

    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 _namespace_for_cluster(cluster_id: str) -> str:
    return f"tidb{cluster_id}"


def _inventory_query(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, *, range_seconds: int, by_instance: bool) -> str:
    selector = (
        f'node_disk_reads_completed_total{{instance=~"{instance_regex}",device!~"dm.*"}}'
    )
    rate_expr = f"rate({selector}[{range_seconds}s])"
    if by_instance:
        return f"sum by (instance) ({rate_expr})"
    return f"sum({rate_expr})"


def _isoformat_utc(ts: int) -> str:
    return datetime.fromtimestamp(ts, tz=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 _vector_value(item: dict[str, Any]) -> float | None:
    raw_value = item.get("value")
    if not isinstance(raw_value, list) or len(raw_value) < 2:
        return None
    return _safe_float(raw_value[1])


def _rank_tikv_nodes_by_instant_value(
    payload: dict[str, Any],
    *,
    node_to_tikv_node: dict[str, str],
) -> list[dict[str, Any]]:
    ranked: list[dict[str, Any]] = []
    for item in _result_series(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"
            )
        avg_value = _vector_value(item)
        if avg_value is None:
            continue
        ranked.append(
            {
                "tikv_node": tikv_node,
                "avg_read_iops": round(avg_value, 3),
            }
        )
    ranked.sort(key=lambda item: (-float(item["avg_read_iops"]), str(item["tikv_node"])))
    return ranked


def _inventory_pairs(payload: dict[str, Any]) -> list[TiKVNodeInventory]:
    resolved: dict[str, str] = {}
    for item in _result_series(payload):
        metric = item.get("metric")
        if not isinstance(metric, dict):
            continue
        node = str(metric.get("node") or "").strip()
        tikv_node = str(metric.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 _resolve_inventory_mapping(
    client: ClinicMetricsClient,
    *,
    cluster_id: str,
    candidate_times: list[int],
) -> tuple[dict[str, str], str]:
    promql = _inventory_query(cluster_id)
    for query_time in candidate_times:
        payload = client.query_instant(
            cluster_id=cluster_id,
            promql=promql,
            query_time=query_time,
        )
        inventory = _inventory_pairs(payload)
        if inventory:
            return {item.node: item.tikv_node for item in inventory}, promql
    raise MetricsQueryError("failed to resolve TiKV node inventory for the target cluster")


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 _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,
    )
    return {
        "id": f"bjt-{local_date}-{args.local_hour:02d}{args.local_minute:02d}-{args.duration_minutes}m",
        "cluster_id": args.cluster_id,
        "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,
        "top_n": args.top_n,
        "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)
    duration_seconds = end_ts - start_ts
    if duration_seconds <= 0:
        raise MetricsQueryError("reference window end must be greater than start")
    midpoint_ts = start_ts + ((end_ts - start_ts) // 2)
    node_to_tikv_node, inventory_query = _resolve_inventory_mapping(
        client,
        cluster_id=args.cluster_id,
        candidate_times=[end_ts, midpoint_ts, start_ts],
    )
    instance_regex = _instance_regex(list(node_to_tikv_node))
    per_instance_query = _disk_read_query(
        instance_regex,
        range_seconds=duration_seconds,
        by_instance=True,
    )
    payload = client.query_instant(
        cluster_id=args.cluster_id,
        promql=per_instance_query,
        query_time=end_ts,
    )
    ranked_nodes = _rank_tikv_nodes_by_instant_value(
        payload,
        node_to_tikv_node=node_to_tikv_node,
    )
    if len(ranked_nodes) < args.top_n:
        raise MetricsQueryError(
            f"expected at least {args.top_n} ranked TiKV nodes, got {len(ranked_nodes)}"
        )
    return {
        "cluster_id": args.cluster_id,
        "window": {
            "start": _isoformat_utc(start_ts),
            "end": _isoformat_utc(end_ts),
        },
        "mapping_strategy": "node_filter",
        "evidence_sufficient": True,
        "node_inventory_count": len(node_to_tikv_node),
        "top_n": args.top_n,
        "top_nodes_by_avg_read_iops": ranked_nodes[: args.top_n],
        "queries": {
            "inventory_query": inventory_query,
            "per_instance_query": per_instance_query,
        },
        "node_to_tikv_node": node_to_tikv_node,
    }


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

    resolve_parser = subparsers.add_parser("resolve-window")
    resolve_parser.add_argument("--cluster-id", 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=300)
    resolve_parser.add_argument("--top-n", type=int, default=3)

    reference_parser = subparsers.add_parser("build-reference")
    reference_parser.add_argument("--cluster-id", 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=300)
    reference_parser.add_argument("--top-n", type=int, default=3)

    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:]))