from __future__ import annotations

import argparse
import json
import math
import os
import sys
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime
from datetime import timedelta
from datetime import timezone
from typing import Any
from zoneinfo import ZoneInfo


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


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_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 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 _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 _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 _vector_value(payload: dict[str, Any]) -> float | None:
    for item in _result_series(payload):
        raw_value = item.get("value")
        if not isinstance(raw_value, list) or len(raw_value) < 2:
            continue
        parsed = _safe_float(raw_value[1])
        if parsed is not None:
            return parsed
    return None


def _escape_label_value(value: str) -> str:
    return value.replace("\\", "\\\\").replace('"', '\\"')


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 _series_match(metric: str, *, db: str, cf: str, level: str | None = None) -> str:
    selector = [f'db="{db}"', f'cf="{cf}"']
    if level is not None:
        selector.append(f'level="{level}"')
    return f'{metric}{{{",".join(selector)}}}'


def _metric_window_max_query(
    metric: str,
    *,
    instance: str,
    db: str,
    cf: str,
    window_seconds: int,
    level: str | None = None,
) -> str:
    selector = [
        f'instance="{_escape_label_value(instance)}"',
        f'db="{db}"',
        f'cf="{cf}"',
    ]
    if level is not None:
        selector.append(f'level="{level}"')
    return f"max(max_over_time({metric}{{{','.join(selector)}}}[{window_seconds}s]))"


def _candidate_score(series_labels: dict[str, Any], *, target_tikv_node: str) -> tuple[int, str]:
    instance = str(series_labels.get("instance") or "").strip()
    if not instance:
        return (0, "")

    for label_name in ("pod", "pod_name", "kubernetes_pod_name"):
        pod_value = str(series_labels.get(label_name) or "").strip()
        if pod_value == target_tikv_node:
            return (3, instance)

    if instance == target_tikv_node:
        return (2, instance)

    if instance.startswith(f"{target_tikv_node}.") or instance.startswith(f"{target_tikv_node}:"):
        return (1, instance)

    return (0, "")


def _resolve_target_instance(
    client: ClinicMetricsClient,
    *,
    cluster_id: str,
    target_tikv_node: str,
    start_ts: int,
    end_ts: int,
) -> tuple[str, dict[str, str]]:
    matches = {
        "pending_compaction": _series_match(
            "tikv_engine_pending_compaction_bytes",
            db="kv",
            cf="write",
        ),
        "l0_files": _series_match(
            "tikv_engine_num_files_at_level",
            db="kv",
            cf="write",
            level="0",
        ),
    }
    best_score = 0
    candidates: set[str] = set()
    for match in matches.values():
        payload = client.query_series(
            cluster_id=cluster_id,
            match=match,
            start=start_ts,
            end=end_ts,
        )
        for item in _series_discovery_items(payload):
            score, instance = _candidate_score(item, target_tikv_node=target_tikv_node)
            if score <= 0 or not instance:
                continue
            if score > best_score:
                best_score = score
                candidates = {instance}
            elif score == best_score:
                candidates.add(instance)

    if not candidates:
        raise MetricsQueryError(
            f"failed to resolve metric instance for target TiKV node '{target_tikv_node}'"
        )
    if len(candidates) != 1:
        raise MetricsQueryError(
            "resolved multiple candidate metric instances for "
            f"'{target_tikv_node}': {', '.join(sorted(candidates))}"
        )
    return next(iter(candidates)), matches


def _rounded_metric_value(payload: dict[str, Any], *, query_name: str) -> int:
    value = _vector_value(payload)
    if value is None:
        raise MetricsQueryError(f"query returned no numeric value for {query_name}")
    return int(round(value))


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


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

    target_instance, discovery_matches = _resolve_target_instance(
        client,
        cluster_id=args.cluster_id,
        target_tikv_node=args.target_tikv_node,
        start_ts=start_ts,
        end_ts=end_ts,
    )

    signal_queries = {
        "write_cf_pending_compaction_bytes_max": _metric_window_max_query(
            "tikv_engine_pending_compaction_bytes",
            instance=target_instance,
            db="kv",
            cf="write",
            window_seconds=window_seconds,
        ),
        "lock_cf_pending_compaction_bytes_max": _metric_window_max_query(
            "tikv_engine_pending_compaction_bytes",
            instance=target_instance,
            db="kv",
            cf="lock",
            window_seconds=window_seconds,
        ),
        "write_cf_l0_files_max": _metric_window_max_query(
            "tikv_engine_num_files_at_level",
            instance=target_instance,
            db="kv",
            cf="write",
            level="0",
            window_seconds=window_seconds,
        ),
        "lock_cf_l0_files_max": _metric_window_max_query(
            "tikv_engine_num_files_at_level",
            instance=target_instance,
            db="kv",
            cf="lock",
            level="0",
            window_seconds=window_seconds,
        ),
    }

    signals: dict[str, int] = {}
    for signal_name, promql in signal_queries.items():
        payload = client.query_instant(
            cluster_id=args.cluster_id,
            promql=promql,
            query_time=end_ts,
        )
        signals[signal_name] = _rounded_metric_value(payload, query_name=signal_name)

    return {
        "cluster_id": args.cluster_id,
        "cluster_label": args.cluster_label,
        "tikv_node": args.target_tikv_node,
        "window": {
            "start": _isoformat_utc(start_ts),
            "end": _isoformat_utc(end_ts),
        },
        "evidence_sufficient": True,
        "signals": signals,
        "resolved_instance": target_instance,
        "queries": {
            "discovery_matches": discovery_matches,
            "signal_queries": signal_queries,
        },
    }


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Runtime hook for single-node TiKV compaction signals")
    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("--target-tikv-node", 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)

    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("--target-tikv-node", required=True)
    reference_parser.add_argument("--start", required=True)
    reference_parser.add_argument("--end", required=True)

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