from __future__ import annotations

import argparse
import json
import math
import os
import re
import sys
import urllib.error
import urllib.parse
import urllib.request
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")
TIDB_NODE_PATTERN = re.compile(r"\bdb(?:-[A-Za-z0-9]+)?-tidb-[A-Za-z0-9-]+\b")
TIKV_NODE_PATTERN = re.compile(r"\bdb(?:-[A-Za-z0-9]+)?-tikv-[A-Za-z0-9-]+\b")
PD_NODE_PATTERN = re.compile(r"\bdb(?:-[A-Za-z0-9]+)?-(?:pd|tso)-[A-Za-z0-9-]+\b")
OUTPUT_PRECISION_DECIMALS = 6
RATE_WINDOW_SECONDS = 120
TIDB_P99_BUCKET_METRIC = "tidb_server_handle_query_duration_seconds_bucket"
BREAKDOWN_METRICS: tuple[tuple[str, str, tuple[str, ...]], ...] = (
    ("avg_plan_build_p99_s", "tidb_session_compile_duration_seconds_bucket", ()),
    ("avg_distributed_execution_p99_s", "tidb_distsql_handle_query_duration_seconds_bucket", ()),
    (
        "avg_kv_request_p99_s",
        "tidb_tikvclient_request_seconds_bucket",
        ('store!="0"', 'scope="false"', 'type="Get"'),
    ),
    (
        "avg_get_tso_wait_p99_s",
        "pd_client_request_handle_requests_duration_seconds_bucket",
        ('type=~"tso.*"',),
    ),
    ("avg_gc_stw_pause_p99_s", "go_gc_pauses_seconds_bucket", ('component="tidb"',)),
)
CLUSTER_SIDE_METRICS: tuple[
    tuple[str, tuple[str, ...], tuple[str, ...], re.Pattern[str]],
    ...,
] = (
    (
        "cluster_avg_tikv_grpc_request_handle_p99_s",
        ("tikv_grpc_msg_duration_seconds_bucket",),
        ('type="kv_get"',),
        TIKV_NODE_PATTERN,
    ),
    (
        "cluster_avg_pd_tso_request_handle_p99_s",
        (
            "pd_server_handle_tso_duration_seconds_bucket",
            "tso_server_handle_tso_duration_seconds_bucket",
        ),
        (),
        PD_NODE_PATTERN,
    ),
)


class ReferenceBuildError(RuntimeError):
    pass


class ClinicClient:
    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 _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 ReferenceBuildError(f"HTTP {exc.code} for {path}: {body}") from exc
        except urllib.error.URLError as exc:
            raise ReferenceBuildError(f"request failed for {path}: {exc}") from exc
        try:
            data = json.loads(payload)
        except json.JSONDecodeError as exc:
            raise ReferenceBuildError(f"invalid JSON response for {path}: {exc}") from exc
        status = str(data.get("status") or "")
        if status and status != "success":
            raise ReferenceBuildError(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 ReferenceBuildError(
        "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 _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 _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 _metric_node_name(metric: dict[str, Any], *, pattern: re.Pattern[str] = TIDB_NODE_PATTERN) -> str:
    for label_name in ("pod", "pod_name", "kubernetes_pod_name"):
        raw_value = str(metric.get(label_name) or "").strip()
        match = pattern.search(raw_value)
        if match:
            return match.group(0)
    for label_name in ("instance", "job", "exported_instance"):
        raw_value = str(metric.get(label_name) or "").strip()
        match = pattern.search(raw_value)
        if match:
            return match.group(0)
    return ""


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


def _mean(values: list[float]) -> float | None:
    if not values:
        return None
    return sum(values) / len(values)


def _instance_regex(nodes: list[str]) -> str:
    escaped = [re.escape(node).replace(r"\-", "-") for node in nodes if node]
    if not escaped:
        return ""
    return f"^({'|'.join(escaped)})($|[-.:].*)"


def _p99_histogram_promql(
    bucket_metric: str,
    *,
    nodes: list[str] | None = None,
    extra_matchers: tuple[str, ...] = (),
) -> str:
    matchers = [matcher for matcher in extra_matchers if matcher]
    if nodes:
        regex = _instance_regex(nodes)
        if regex:
            matchers.append(f'instance=~"{regex}"')
    selector = f"{{{','.join(matchers)}}}" if matchers else ""
    return (
        "histogram_quantile(0.99, "
        f"sum by (le, instance) (rate({bucket_metric}{selector}[{RATE_WINDOW_SECONDS}s])))"
    )


def _avg_metric_by_named_node(
    payload: dict[str, Any],
    *,
    pattern: re.Pattern[str] = TIDB_NODE_PATTERN,
) -> dict[str, float]:
    values_by_node: dict[str, list[float]] = {}
    for item in _result_series(payload):
        metric = item.get("metric")
        if not isinstance(metric, dict):
            continue
        node = _metric_node_name(metric, pattern=pattern)
        if not node:
            continue
        values = _series_numeric_values(item)
        if not values:
            continue
        values_by_node.setdefault(node, []).extend(values)

    averages: dict[str, float] = {}
    for node, values in values_by_node.items():
        avg_value = _mean(values)
        if avg_value is None:
            continue
        averages[node] = avg_value
    return averages


def _avg_metric_by_tidb_node(payload: dict[str, Any]) -> dict[str, float]:
    return _avg_metric_by_named_node(payload, pattern=TIDB_NODE_PATTERN)


def _cluster_mean(values_by_node: dict[str, float]) -> float | None:
    if not values_by_node:
        return None
    return _mean(list(values_by_node.values()))

def _ranked_node_value_pairs(values_by_node: dict[str, float]) -> list[tuple[str, float]]:
    ranked_pairs = list(values_by_node.items())
    ranked_pairs.sort(key=lambda item: (-float(item[1]), str(item[0])))
    return ranked_pairs


def _rank_tidb_nodes_by_avg_p99(payload: dict[str, Any]) -> list[dict[str, Any]]:
    ranked_pairs = _ranked_node_value_pairs(_avg_metric_by_tidb_node(payload))
    return [
        {
            "node": node,
            "avg_p99_latency_s": round(avg_latency, OUTPUT_PRECISION_DECIMALS),
        }
        for node, avg_latency in ranked_pairs
    ]


def _build_top_node_breakdown(
    ranked_pairs: list[tuple[str, float]],
    metric_values_by_field: dict[str, dict[str, float]],
) -> list[dict[str, Any]]:
    nodes: list[dict[str, Any]] = []
    for node, avg_latency in ranked_pairs:
        row: dict[str, Any] = {
            "node": node,
            "avg_p99_latency_s": round(avg_latency, OUTPUT_PRECISION_DECIMALS),
        }
        for field_name, values_by_node in metric_values_by_field.items():
            row[field_name] = round(values_by_node[node], OUTPUT_PRECISION_DECIMALS)
        nodes.append(row)
    return nodes


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,
        "cluster_label": args.cluster_label,
        "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,
    }


def _build_reference(args: argparse.Namespace) -> dict[str, Any]:
    client = ClinicClient(
        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 ReferenceBuildError("reference window end must be greater than start")
    shared_metrics_cluster_id = args.shared_metrics_cluster_id or args.cluster_id
    payload = client.query_range(
        cluster_id=args.cluster_id,
        promql=_p99_histogram_promql(TIDB_P99_BUCKET_METRIC),
        start=start_ts,
        end=end_ts,
        step_seconds=args.step_seconds,
    )
    primary_values_by_node = _avg_metric_by_tidb_node(payload)
    ranked_pairs = _ranked_node_value_pairs(primary_values_by_node)
    top_ranked_pairs = ranked_pairs[: args.top_n]
    top_nodes = [node for node, _ in top_ranked_pairs]

    metric_values_by_field: dict[str, dict[str, float]] = {}
    for field_name, bucket_metric, extra_matchers in BREAKDOWN_METRICS:
        metric_payload = client.query_range(
            cluster_id=args.cluster_id,
            promql=_p99_histogram_promql(
                bucket_metric,
                nodes=top_nodes,
                extra_matchers=extra_matchers,
            ),
            start=start_ts,
            end=end_ts,
            step_seconds=args.step_seconds,
        )
        metric_values_by_field[field_name] = _avg_metric_by_tidb_node(metric_payload)

    cluster_metric_values: dict[str, float | None] = {}
    for field_name, bucket_metrics, extra_matchers, pattern in CLUSTER_SIDE_METRICS:
        cluster_metric_values[field_name] = None
        for bucket_metric in bucket_metrics:
            metric_payload = client.query_range(
                cluster_id=shared_metrics_cluster_id,
                promql=_p99_histogram_promql(
                    bucket_metric,
                    extra_matchers=extra_matchers,
                ),
                start=start_ts,
                end=end_ts,
                step_seconds=args.step_seconds,
            )
            value = _cluster_mean(
                _avg_metric_by_named_node(metric_payload, pattern=pattern)
            )
            if value is not None:
                cluster_metric_values[field_name] = value
                break

    evidence_sufficient = len(top_ranked_pairs) >= args.top_n and all(
        all(node in metric_values_by_field[field_name] for field_name, _, _ in BREAKDOWN_METRICS)
        for node, _ in top_ranked_pairs
    ) and all(cluster_metric_values[field_name] is not None for field_name, _, _, _ in CLUSTER_SIDE_METRICS)

    result = {
        "cluster_id": args.cluster_id,
        "window": {
            "start": _isoformat_utc(start_ts),
            "end": _isoformat_utc(end_ts),
        },
        "evidence_sufficient": evidence_sufficient,
        "top_nodes_by_avg_p99_latency": (
            _build_top_node_breakdown(top_ranked_pairs, metric_values_by_field)
            if evidence_sufficient
            else []
        ),
    }
    for field_name, value in cluster_metric_values.items():
        result[field_name] = (
            round(value, OUTPUT_PRECISION_DECIMALS) if value is not None else None
        )
    return result


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Runtime hook for TiDB P99 latency ranking")
    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("--cluster-label", 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)
    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("--cluster-label", required=True)
    reference_parser.add_argument("--shared-metrics-cluster-id")
    reference_parser.add_argument("--start", required=True)
    reference_parser.add_argument("--end", required=True)
    reference_parser.add_argument("--step-seconds", type=int, default=60)
    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 ReferenceBuildError(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:]))