from __future__ import annotations

import argparse
import hashlib
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
BEIJING_TZ = ZoneInfo("Asia/Shanghai")
DEFAULT_KEYSPACE_LABEL = "keyspace_name"


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_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, 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=timezone.utc).isoformat().replace("+00:00", "Z")


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(timezone.utc).timestamp()),
        int(candidate_end.astimezone(timezone.utc).timestamp()),
        candidate_start.date().isoformat(),
    )


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 _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 _series_label_keys(series: list[dict[str, Any]]) -> set[str]:
    keys: set[str] = set()
    for item in series:
        metric = item.get("metric")
        if not isinstance(metric, dict):
            continue
        for key in metric.keys():
            if isinstance(key, str):
                keys.add(key)
    return keys


def _series_scope_summary(series: list[dict[str, Any]]) -> dict[str, Any]:
    label_keys = _series_label_keys(series)
    return {
        "series_count": len(series),
        "label_keys": sorted(label_keys),
        "has_keyspace_name": "keyspace_name" in label_keys,
        "has_keyspace_id": "keyspace_id" in label_keys,
        "has_region_id": "region_id" in label_keys,
    }


def _label_values_from_series(series: list[dict[str, Any]], label_name: str) -> list[str]:
    values = {
        str(metric.get(label_name) or "").strip()
        for item in series
        if isinstance((metric := item.get("metric")), dict)
    }
    return sorted(value for value in values if value)


def _patch_workspace_evidence_report(*, start_ts: int, end_ts: int, cluster_id: str) -> None:
    audit_dir = os.environ.get("INVESTIBOOK_EVAL_AUDIT_DIR", "").strip()
    if not audit_dir:
        return
    report_path = os.path.join(audit_dir, "evidences", "workspace_evidence_report.json")
    if not os.path.exists(report_path):
        return
    try:
        with open(report_path, "r", encoding="utf-8") as handle:
            report = json.load(handle)
    except (OSError, json.JSONDecodeError):
        return
    if not isinstance(report, dict):
        return
    metadata = report.setdefault("metadata", {})
    if not isinstance(metadata, dict):
        metadata = {}
        report["metadata"] = metadata
    metadata["cluster_id"] = cluster_id
    metadata["time_range"] = {
        "start": str(start_ts),
        "end": str(end_ts),
    }
    try:
        with open(report_path, "w", encoding="utf-8") as handle:
            json.dump(report, handle, ensure_ascii=False, indent=2)
            handle.write("\n")
    except OSError:
        return


def _instance_from_metric(metric: dict[str, Any]) -> str:
    for label_name in ("instance", "pod", "kubernetes_pod_name"):
        value = str(metric.get(label_name) or "").strip()
        if value:
            return value
    return ""


def _inventory_promql(*, tidb_cluster_id: str | None = None) -> str:
    labels = ['component="tikv"']
    if tidb_cluster_id:
        labels.append(f'tidb_cluster_id="{tidb_cluster_id}"')
    selector = "process_cpu_seconds_total{" + ",".join(labels) + "}"
    return f"sum by (instance, pod) (rate({selector}[1m]))"


def _count_promql(selector: str, *, group_by: tuple[str, ...] = ()) -> str:
    if not group_by:
        return f"count({selector})"
    return f"count by ({', '.join(group_by)}) ({selector})"


def _store_size_selector(
    instance: str | None = None,
    *,
    tidb_cluster_id: str | None = None,
) -> str:
    labels = ['type="used"']
    if instance:
        escaped_instance = instance.replace("\\", "\\\\").replace('"', '\\"')
        labels.append(f'instance="{escaped_instance}"')
    if tidb_cluster_id:
        labels.append(f'tidb_cluster_id="{tidb_cluster_id}"')
    return "tikv_store_size_bytes{" + ",".join(labels) + "}"


def _store_size_promql(instance: str, *, tidb_cluster_id: str | None = None) -> str:
    return _store_size_selector(instance, tidb_cluster_id=tidb_cluster_id)


def _select_target_instance(
    series: list[dict[str, Any]],
    *,
    cluster_id: str,
    start_ts: int,
    end_ts: int,
) -> str:
    instances = sorted(
        {
            _instance_from_metric(item.get("metric", {}))
            for item in series
            if isinstance(item.get("metric"), dict)
        }
    )
    instances = [item for item in instances if item]
    if not instances:
        raise MetricsQueryError("failed to resolve TiKV instances for the target cluster")
    seed = f"{cluster_id}:{start_ts}:{end_ts}:premium-keyspace-storage-query-semantics"
    digest = hashlib.sha256(seed.encode("utf-8")).digest()
    index = int.from_bytes(digest[:8], "big") % len(instances)
    return instances[index]


def _resolve_window(args: argparse.Namespace) -> dict[str, Any]:
    current_start, current_end, local_date = _latest_completed_beijing_window(
        local_hour=args.local_hour,
        local_minute=args.local_minute,
        duration_minutes=args.duration_minutes,
    )
    client = ClinicMetricsClient(
        api_key=_env_api_key(),
        base_url=os.environ.get("CLINIC_ENDPOINT", DEFAULT_BASE_URL).strip() or DEFAULT_BASE_URL,
    )
    inventory_payload = client.query_instant(
        cluster_id=args.metrics_cluster_id,
        promql=_inventory_promql(tidb_cluster_id=args.metrics_cluster_id),
        query_time=current_end,
    )
    instances = _result_series(inventory_payload)
    target_instance = _select_target_instance(
        instances,
        cluster_id=args.cluster_id,
        start_ts=current_start,
        end_ts=current_end,
    )
    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,
        "metrics_cluster_id": args.metrics_cluster_id,
        "local_timezone": "Asia/Shanghai",
        "local_date": local_date,
        "start": _isoformat_utc(current_start),
        "end": _isoformat_utc(current_end),
        "start_ts": current_start,
        "end_ts": current_end,
        "target_tikv_instance": target_instance,
    }


def _build_cluster_result(
    *,
    cluster_id: str,
    metrics_cluster_id: str,
    target_instance: str,
    keyspace_label: str,
    client: ClinicMetricsClient,
    start_ts: int,
    end_ts: int,
    step_seconds: int,
) -> dict[str, Any]:
    selector = _store_size_selector(tidb_cluster_id=metrics_cluster_id)
    payload = client.query_instant(
        cluster_id=cluster_id,
        promql=selector,
        query_time=end_ts,
    )
    series = _result_series(payload)
    keyspace_values = _label_values_from_series(series, keyspace_label)
    return {
        "cluster_id": cluster_id,
        "target_tikv_instance": target_instance,
        "keyspace_label": keyspace_label,
        "keyspace_value_count": len(keyspace_values),
    }


def _build_reference(args: argparse.Namespace) -> dict[str, Any]:
    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")

    client = ClinicMetricsClient(
        api_key=_env_api_key(),
        base_url=os.environ.get("CLINIC_ENDPOINT", DEFAULT_BASE_URL).strip() or DEFAULT_BASE_URL,
    )
    cluster_result = _build_cluster_result(
        cluster_id=args.cluster_id,
        metrics_cluster_id=args.metrics_cluster_id,
        target_instance=args.target_tikv_instance,
        keyspace_label=args.keyspace_label,
        client=client,
        start_ts=start_ts,
        end_ts=end_ts,
        step_seconds=args.step_seconds,
    )
    _patch_workspace_evidence_report(
        start_ts=start_ts,
        end_ts=end_ts,
        cluster_id=args.cluster_id,
    )
    evidence_sufficient = bool(cluster_result["keyspace_value_count"])
    return {
        "target_metric": "tikv_store_size_bytes",
        "evidence_sufficient": evidence_sufficient,
        "window": {
            "start": args.start,
            "end": args.end,
        },
        "clusters": [cluster_result] if evidence_sufficient else [],
    }


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Runtime hooks for premium-keyspace-storage-query-semantics")
    subparsers = parser.add_subparsers(dest="command", required=True)

    resolve_window = subparsers.add_parser("resolve-window", help="Resolve the latest completed Beijing window")
    resolve_window.add_argument("--cluster-id", required=True)
    resolve_window.add_argument("--cluster-label", required=True)
    resolve_window.add_argument("--metrics-cluster-id", required=True)
    resolve_window.add_argument("--local-hour", type=int, default=11)
    resolve_window.add_argument("--local-minute", type=int, default=0)
    resolve_window.add_argument("--duration-minutes", type=int, default=5)
    resolve_window.add_argument("--step-seconds", type=int, default=60)
    resolve_window.set_defaults(handler=_resolve_window)

    build_reference = subparsers.add_parser("build-reference", help="Build the runtime reference answer")
    build_reference.add_argument("--cluster-id", required=True)
    build_reference.add_argument("--metrics-cluster-id", required=True)
    build_reference.add_argument("--target-tikv-instance", required=True)
    build_reference.add_argument("--keyspace-label", required=True)
    build_reference.add_argument("--start", required=True)
    build_reference.add_argument("--end", required=True)
    build_reference.add_argument("--step-seconds", type=int, default=60)
    build_reference.set_defaults(handler=_build_reference)

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)
    output = args.handler(args)
    json.dump(output, sys.stdout, ensure_ascii=False, indent=2)
    sys.stdout.write("\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())