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")
COMPONENT_ORDER = ("tidb", "tikv", "pd")
CPU_PROMQL_BY_COMPONENT = {
    "tidb": 'sum by (pod, instance) (rate(process_cpu_seconds_total{component="tidb"}[1m]))',
    "tikv": 'sum by (pod, instance) (rate(process_cpu_seconds_total{component="tikv"}[1m]))',
    "pd": 'sum by (pod, instance) (rate(process_cpu_seconds_total{component="pd"}[1m]))',
}
QUOTA_PROMQL_BY_COMPONENT = {
    "tidb": "max by (pod, instance) (tidb_server_maxprocs)",
    "tikv": "max by (pod, instance) (tikv_server_cpu_cores_quota)",
    "pd": "max by (pod, instance) (pd_service_maxprocs)",
}
NODE_PATTERN_BY_COMPONENT = {
    "tidb": re.compile(r"\bdb-tidb-[A-Za-z0-9-]+\b"),
    "tikv": re.compile(r"\bdb-tikv-[A-Za-z0-9-]+\b"),
    "pd": re.compile(r"\bdb-pd-[A-Za-z0-9-]+\b"),
}
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 loki_query_range(
        self,
        *,
        cluster_id: str,
        logql: str,
        start_ns: int,
        end_ns: int,
        limit: int,
        direction: str = "backward",
    ) -> dict[str, Any]:
        return self._get(
            f"/data-proxy/loki/{cluster_id}/api/v1/query_range",
            params={
                "query": logql,
                "start": str(start_ns),
                "end": str(end_ns),
                "limit": str(limit),
                "direction": direction,
            },
        )

    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 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 _error_warn_level(line: str) -> str:
    upper = str(line or "").upper()
    if "[ERROR]" in upper or "[FATAL]" in upper:
        return "ERROR"
    if "[WARN]" in upper or "[WARNING]" in upper:
        return "WARN"
    return ""


def _component_label_matcher(component: str) -> str:
    if component == "tidb":
        return "tidb|statementlog"
    return component


def _metric_node_name(metric: dict[str, Any], *, component: str) -> str:
    for label_name in ("pod", "pod_name", "kubernetes_pod_name"):
        pod = str(metric.get(label_name) or "").strip()
        if pod.startswith(f"db-{component}-"):
            return pod

    matcher = NODE_PATTERN_BY_COMPONENT[component]
    for label_name in ("instance", "job", "exported_instance", "pod", "pod_name", "kubernetes_pod_name"):
        raw_value = str(metric.get(label_name) or "").strip()
        match = matcher.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 _median(values: list[float]) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    return ordered[len(ordered) // 2]


def _cpu_overload_rows(
    client: ClinicClient,
    *,
    cluster_id: str,
    component: str,
    start_ts: int,
    end_ts: int,
    top_n: int,
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    cpu_query = CPU_PROMQL_BY_COMPONENT[component]
    quota_query = QUOTA_PROMQL_BY_COMPONENT[component]
    cpu_payload = client.query_range(
        cluster_id=cluster_id,
        promql=cpu_query,
        start=start_ts,
        end=end_ts,
        step_seconds=60,
    )
    quota_payload = client.query_range(
        cluster_id=cluster_id,
        promql=quota_query,
        start=start_ts,
        end=end_ts,
        step_seconds=60,
    )

    cpu_by_node: dict[str, float] = {}
    for item in _result_series(cpu_payload):
        metric = item.get("metric")
        if not isinstance(metric, dict):
            continue
        node = _metric_node_name(metric, component=component)
        if not node:
            continue
        mean_cpu = _mean(_series_numeric_values(item))
        if mean_cpu is None:
            continue
        cpu_by_node[node] = max(cpu_by_node.get(node, 0.0), mean_cpu)

    quota_by_node: dict[str, float] = {}
    for item in _result_series(quota_payload):
        metric = item.get("metric")
        if not isinstance(metric, dict):
            continue
        node = _metric_node_name(metric, component=component)
        if not node:
            continue
        quota = _median(_series_numeric_values(item))
        if quota is None or quota <= 0:
            continue
        quota_by_node[node] = max(quota_by_node.get(node, 0.0), quota)

    rows: list[dict[str, Any]] = []
    for node, mean_cpu in cpu_by_node.items():
        quota = quota_by_node.get(node)
        if quota is None or quota <= 0:
            continue
        overload = mean_cpu / quota
        rows.append(
            {
                "component": component,
                "node": node,
                "avg_cpu_overload": round(overload, 4),
                "avg_cpu_cores": round(mean_cpu, 4),
                "cpu_quota": round(quota, 4),
            }
        )

    rows.sort(key=lambda item: (-float(item["avg_cpu_overload"]), str(item["node"])))
    selected = rows[:top_n]
    for index, row in enumerate(selected, start=1):
        row["rank"] = index
    return selected, {
        "cpu_query": cpu_query,
        "quota_query": quota_query,
        "observed_cpu_nodes": len(cpu_by_node),
        "observed_quota_nodes": len(quota_by_node),
        "selected_count": len(selected),
    }


def _loki_time_ns(value: int) -> int:
    return int(value) * 1_000_000_000


def _loki_streams(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 _stream_node_name(labels: dict[str, Any], line: str, *, component: str) -> str:
    matcher = NODE_PATTERN_BY_COMPONENT[component]
    for label_name in ("pod", "instance", "app", "service_name", "filename", "job", "container", "component"):
        raw_value = str(labels.get(label_name) or "").strip()
        match = matcher.search(raw_value)
        if match:
            return match.group(0)
    return _metric_node_name({"instance": line}, component=component)


def _node_log_queries(node: str) -> list[tuple[str, str, str]]:
    filter_expr = r"|~ `(?i)\[error\]|\[warn(?:ing)?\]`"
    return [
        (
            "instance_selector",
            f'{{instance="{node}"}} {filter_expr}',
            "1000",
        ),
        (
            "app_selector",
            f'{{app="{node}"}} {filter_expr}',
            "1000",
        ),
        (
            "service_name_selector",
            f'{{service_name="{node}"}} {filter_expr}',
            "1000",
        ),
    ]


def _time_slices(start_ts: int, end_ts: int, *, step_seconds: int = 60) -> list[tuple[int, int]]:
    slices: list[tuple[int, int]] = []
    cursor = start_ts
    while cursor < end_ts:
        next_cursor = min(end_ts, cursor + step_seconds)
        slices.append((cursor, next_cursor))
        cursor = next_cursor
    return slices


def _aggregate_component_logs(
    client: ClinicClient,
    *,
    cluster_id: str,
    component: str,
    target_nodes: set[str],
    start_ts: int,
    end_ts: int,
) -> tuple[dict[str, int], dict[str, Any]]:
    counts_by_node: dict[str, int] = {node: 0 for node in target_nodes}
    seen_lines: set[tuple[str, str, str]] = set()
    debug_queries: list[dict[str, Any]] = []

    for node in sorted(target_nodes):
        for strategy, logql, limit in _node_log_queries(node):
            matched_lines = 0
            for slice_start, slice_end in _time_slices(start_ts, end_ts):
                payload = client.loki_query_range(
                    cluster_id=cluster_id,
                    logql=logql,
                    start_ns=_loki_time_ns(slice_start),
                    end_ns=_loki_time_ns(slice_end),
                    limit=int(limit),
                )
                for stream in _loki_streams(payload):
                    labels = stream.get("stream")
                    if not isinstance(labels, dict):
                        labels = {}
                    for entry in stream.get("values") or []:
                        if not isinstance(entry, list) or len(entry) < 2:
                            continue
                        ts_ns = str(entry[0] or "")
                        line = str(entry[1] or "")
                        matched_node = _stream_node_name(labels, line, component=component)
                        if matched_node != node:
                            continue
                        level = _error_warn_level(line)
                        if level not in {"ERROR", "WARN"}:
                            continue
                        dedupe_key = (ts_ns, node, line)
                        if dedupe_key in seen_lines:
                            continue
                        seen_lines.add(dedupe_key)
                        counts_by_node[node] = counts_by_node.get(node, 0) + 1
                        matched_lines += 1
            debug_queries.append(
                {
                    "node": node,
                    "strategy": strategy,
                    "logql": logql,
                    "matched_lines": matched_lines,
                }
            )
    return counts_by_node, {"queries": debug_queries}


def _resolve_window(args: argparse.Namespace) -> dict[str, Any]:
    start_ts, end_ts, 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(start_ts),
        "end": _isoformat_utc(end_ts),
        "local_timezone": "Asia/Shanghai",
        "local_date": local_date,
        "duration_minutes": args.duration_minutes,
        "top_n_per_component": args.top_n_per_component,
        "cpu_rule": "mean(rate(process_cpu_seconds_total[1m])) divided by component quota/maxprocs",
        "log_rule": "error_warn_only_unique_line_count_per_node",
    }


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 ReferenceBuildError("reference window end must be greater than start")

    client = ClinicClient(
        api_key=_env_api_key(),
        base_url=os.environ.get("CLINIC_ENDPOINT", DEFAULT_BASE_URL).strip() or DEFAULT_BASE_URL,
        timeout_s=args.timeout_s,
    )

    ranked_nodes: list[dict[str, Any]] = []
    debug_components: dict[str, Any] = {}

    for component in COMPONENT_ORDER:
        selected_rows, component_debug = _cpu_overload_rows(
            client,
            cluster_id=args.cluster_id,
            component=component,
            start_ts=start_ts,
            end_ts=end_ts,
            top_n=args.top_n_per_component,
        )
        debug_components[component] = component_debug
        if len(selected_rows) < args.top_n_per_component:
            return {
                "cluster_id": args.cluster_id,
                "window": {
                    "start": _isoformat_utc(start_ts),
                    "end": _isoformat_utc(end_ts),
                },
                "evidence_sufficient": False,
                "ranked_nodes": [],
                "debug": {
                    "components": debug_components,
                    "failure_component": component,
                },
            }

        target_nodes = {str(row["node"]) for row in selected_rows}
        log_counts_by_node, logs_debug = _aggregate_component_logs(
            client,
            cluster_id=args.cluster_id,
            component=component,
            target_nodes=target_nodes,
            start_ts=start_ts,
            end_ts=end_ts,
        )
        component_debug["log_queries"] = logs_debug["queries"]

        for row in selected_rows:
            ranked_nodes.append(
                {
                    "component": component,
                    "rank": int(row["rank"]),
                    "node": str(row["node"]),
                    "avg_cpu_overload": float(row["avg_cpu_overload"]),
                    "error_warn_log_count": int(log_counts_by_node.get(str(row["node"]), 0)),
                }
            )

    ranked_nodes.sort(
        key=lambda item: (
            COMPONENT_ORDER.index(str(item["component"])),
            int(item["rank"]),
        )
    )
    return {
        "cluster_id": args.cluster_id,
        "window": {
            "start": _isoformat_utc(start_ts),
            "end": _isoformat_utc(end_ts),
        },
        "evidence_sufficient": True,
        "ranked_nodes": ranked_nodes,
        "debug": {
            "components": debug_components,
        },
    }


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Runtime hooks for top-cpu-overload-nodes-and-log-patterns-in-saas01-during-latest-window"
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    resolve_window = subparsers.add_parser(
        "resolve-window",
        help="Resolve the latest completed Beijing 11:00-11:05 window",
    )
    resolve_window.add_argument("--cluster-id", required=True)
    resolve_window.add_argument("--cluster-label", required=False)
    resolve_window.add_argument("--local-hour", type=int, required=True)
    resolve_window.add_argument("--local-minute", type=int, required=True)
    resolve_window.add_argument("--duration-minutes", type=int, required=True)
    resolve_window.add_argument("--top-n-per-component", type=int, required=True)
    resolve_window.set_defaults(handler=_resolve_window)

    build_reference = subparsers.add_parser(
        "build-reference",
        help="Build the live CPU-overload and error/warn log-count reference for the resolved window",
    )
    build_reference.add_argument("--cluster-id", required=True)
    build_reference.add_argument("--cluster-label", required=False)
    build_reference.add_argument("--start", required=True)
    build_reference.add_argument("--end", required=True)
    build_reference.add_argument("--top-n-per-component", type=int, required=True)
    build_reference.add_argument("--timeout-s", type=int, default=DEFAULT_TIMEOUT_S)
    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)
    try:
        output = args.handler(args)
    except ReferenceBuildError as exc:
        print(str(exc), file=sys.stderr)
        return 1
    json.dump(output, sys.stdout, ensure_ascii=False, indent=2)
    sys.stdout.write("\n")
    return 0


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