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 pathlib import Path
from typing import Any
from zoneinfo import ZoneInfo


DEFAULT_TIMEOUT_S = 60
DEFAULT_O11Y_AUTH_URL = "https://www.ds.us-west-2.aws.observability.tidbcloud.com/client/auth"
BEIJING_TZ = ZoneInfo("Asia/Shanghai")


@dataclass(frozen=True)
class TiKVNodeInventory:
    node: str
    tikv_node: str


class MetricsQueryError(RuntimeError):
    pass


class O11yMetricsClient:
    def __init__(
        self,
        *,
        provider: str,
        region: str,
        infra_uuid: str,
        github_pat: str,
        timeout_s: int = DEFAULT_TIMEOUT_S,
    ) -> None:
        self._provider = provider
        self._region = region
        self._infra_uuid = infra_uuid
        self._timeout_s = timeout_s
        self._auth0_token = self._authenticate(github_pat=github_pat)
        self._metrics_base = (
            "https://www.ds."
            f"{self._region}.{self._provider}.observability.tidbcloud.com/internal/v2/metrics/"
            f"{self._infra_uuid}/api/v1"
        )

    def query_range(
        self,
        *,
        promql: str,
        start: int,
        end: int,
        step_seconds: int,
    ) -> dict[str, Any]:
        return self._get(
            "/query_range",
            params={
                "query": promql,
                "start": str(start),
                "end": str(end),
                "step": str(step_seconds),
            },
        )

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

    def _authenticate(self, *, github_pat: str) -> str:
        body = json.dumps(
            {
                "github_token": github_pat,
                "provider": self._provider,
            }
        ).encode("utf-8")
        request = urllib.request.Request(
            DEFAULT_O11Y_AUTH_URL,
            headers={"Content-Type": "application/json"},
            data=body,
            method="POST",
        )
        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} during O11Y auth: {body}") from exc
        except urllib.error.URLError as exc:
            raise MetricsQueryError(f"O11Y auth request failed: {exc}") from exc
        try:
            data = json.loads(payload)
        except json.JSONDecodeError as exc:
            raise MetricsQueryError(f"invalid O11Y auth JSON response: {exc}") from exc
        auth0_token = str(data.get("auth0_token") or "").strip()
        if not auth0_token:
            raise MetricsQueryError("missing auth0_token in O11Y auth response")
        return auth0_token

    def _get(self, path: str, *, params: dict[str, str]) -> dict[str, Any]:
        url = f"{self._metrics_base}{path}?{urllib.parse.urlencode(params)}"
        request = urllib.request.Request(
            url,
            headers={"Authorization": f"Bearer {self._auth0_token}"},
            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_github_token() -> str:
    for key in (
        "SKILL_GITHUB_TOKEN",
        "INVESTIBOOK_GITHUB_TOKEN",
        "GITHUB_TOKEN",
        "GH_TOKEN",
        "GITHUB_PAT",
    ):
        value = os.environ.get(key, "").strip()
        if value:
            return value
    raise MetricsQueryError(
        "missing GitHub token in "
        "SKILL_GITHUB_TOKEN/INVESTIBOOK_GITHUB_TOKEN/GITHUB_TOKEN/GH_TOKEN/GITHUB_PAT"
    )


def _inventory_query(namespace: str) -> str:
    return f'kube_pod_info{{namespace="{namespace}",created_by_kind="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 _load_json_object_from_path(path: Path) -> dict[str, Any]:
    try:
        raw = path.read_text(encoding="utf-8")
    except OSError as exc:
        raise MetricsQueryError(f"failed to read JSON file '{path}': {exc}") from exc
    try:
        payload = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise MetricsQueryError(f"invalid JSON file '{path}': {exc}") from exc
    if not isinstance(payload, dict):
        raise MetricsQueryError(f"JSON file '{path}' must contain an object")
    return payload


def _resolve_case_relative_path(raw_path: str) -> Path:
    path = Path(raw_path).expanduser()
    if path.is_absolute():
        return path
    return Path(__file__).resolve().parent / path


def _prefetched_ranked_nodes(payload: dict[str, Any]) -> tuple[list[dict[str, Any]], dict[str, str]]:
    direct_rows = payload.get("top_nodes_by_avg_read_iops")
    if isinstance(direct_rows, list):
        ranked: list[dict[str, Any]] = []
        for item in direct_rows:
            if not isinstance(item, dict):
                raise MetricsQueryError("prefetched top_nodes_by_avg_read_iops entries must be objects")
            tikv_node = str(item.get("tikv_node") or "").strip()
            avg_value = _safe_float(item.get("avg_read_iops"))
            if not tikv_node or avg_value is None:
                raise MetricsQueryError(
                    "prefetched top_nodes_by_avg_read_iops entries require tikv_node and avg_read_iops"
                )
            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, {}

    inventory_raw = payload.get("inventory")
    metrics_raw = payload.get("per_instance_read_iops")
    if not isinstance(inventory_raw, list) or not isinstance(metrics_raw, list):
        raise MetricsQueryError(
            "prefetched input requires either top_nodes_by_avg_read_iops or both inventory and per_instance_read_iops"
        )

    instance_to_tikv_node: dict[str, str] = {}
    for item in inventory_raw:
        if not isinstance(item, dict):
            raise MetricsQueryError("prefetched inventory entries must be objects")
        instance = str(item.get("instance") or item.get("node") or "").strip()
        tikv_node = str(item.get("tikv_node") or item.get("pod") or "").strip()
        if not instance or not tikv_node:
            raise MetricsQueryError("prefetched inventory entries require instance and tikv_node")
        previous = instance_to_tikv_node.get(instance)
        if previous and previous != tikv_node:
            raise MetricsQueryError(
                f"prefetched inventory maps instance '{instance}' to multiple TiKV nodes"
            )
        instance_to_tikv_node[instance] = tikv_node

    ranked: list[dict[str, Any]] = []
    for item in metrics_raw:
        if not isinstance(item, dict):
            raise MetricsQueryError("prefetched per_instance_read_iops entries must be objects")
        instance = str(item.get("instance") or "").strip()
        avg_value = _safe_float(item.get("avg_read_iops"))
        if not instance or avg_value is None:
            raise MetricsQueryError("prefetched per_instance_read_iops entries require instance and avg_read_iops")
        tikv_node = instance_to_tikv_node.get(instance)
        if not tikv_node:
            raise MetricsQueryError(
                f"prefetched metric instance '{instance}' was not found in prefetched inventory"
            )
        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, instance_to_tikv_node


def _validate_prefetched_window(
    payload: dict[str, Any],
    *,
    expected_start: str,
    expected_end: str,
) -> None:
    window = payload.get("window")
    if not isinstance(window, dict):
        return
    start = str(window.get("start") or "").strip()
    end = str(window.get("end") or "").strip()
    if start and start != expected_start:
        raise MetricsQueryError(
            f"prefetched input window.start '{start}' does not match expected '{expected_start}'"
        )
    if end and end != expected_end:
        raise MetricsQueryError(
            f"prefetched input window.end '{end}' does not match expected '{expected_end}'"
        )


def _workspace_dir() -> Path:
    raw = os.environ.get("INVESTIBOOK_EVAL_WORKSPACE_DIR", "").strip()
    if not raw:
        raise MetricsQueryError("missing INVESTIBOOK_EVAL_WORKSPACE_DIR")
    return Path(raw)


def _runtime_context_path() -> Path:
    raw = os.environ.get("INVESTIBOOK_EVAL_RUNTIME_CONTEXT_PATH", "").strip()
    if not raw:
        raise MetricsQueryError("missing INVESTIBOOK_EVAL_RUNTIME_CONTEXT_PATH")
    return Path(raw)


def _write_workspace_prefetched_artifact(*, payload: dict[str, Any], filename: str) -> str:
    workspace_dir = _workspace_dir()
    evidence_dir = workspace_dir / "evidences"
    evidence_dir.mkdir(parents=True, exist_ok=True)
    artifact_path = evidence_dir / filename
    artifact_path.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return str(artifact_path.relative_to(workspace_dir))


def _workspace_path_from_relative(relative_path: str) -> Path:
    path = Path(relative_path)
    if path.is_absolute():
        return path
    return _workspace_dir() / path


def _load_runtime_output_data(output_id: str) -> dict[str, Any]:
    runtime_context = _load_json_object_from_path(_runtime_context_path())
    outputs = runtime_context.get("outputs")
    if not isinstance(outputs, dict):
        raise MetricsQueryError("runtime context is missing outputs")
    item = outputs.get(output_id)
    if not isinstance(item, dict):
        raise MetricsQueryError(f"runtime context is missing output '{output_id}'")
    data = item.get("data")
    if not isinstance(data, dict):
        raise MetricsQueryError(f"runtime output '{output_id}' does not contain an object payload")
    return data


def _prepare_prefetched_case_input(
    *,
    prefetched_input_path: str,
    artifact_filename: str,
    expected_start: str,
    expected_end: str,
) -> tuple[str, str]:
    source_path = _resolve_case_relative_path(prefetched_input_path)
    payload = _load_json_object_from_path(source_path)
    _validate_prefetched_window(
        payload,
        expected_start=expected_start,
        expected_end=expected_end,
    )
    _prefetched_ranked_nodes(payload)
    relative_path = _write_workspace_prefetched_artifact(
        payload=payload,
        filename=artifact_filename,
    )
    prompt = (
        f"workspace 中已提供预取好的基础设施证据文件 `{relative_path}`。"
        "请优先使用该文件完成排序，不要重复查询同类 `node_*` / `kube_*` / `aws_ebs_*` 指标。"
    )
    return relative_path, prompt


def _resolve_inventory_mapping(
    client: O11yMetricsClient,
    *,
    namespace: str,
    candidate_times: list[int],
) -> tuple[dict[str, str], str]:
    promql = _inventory_query(namespace)
    for query_time in candidate_times:
        payload = client.query_instant(
            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 premium 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 _prepare_case_input(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,
    )
    prefetched_relative_path = ""
    prefetched_prompt = ""
    evidence_mode = "live_query"
    prefetched_input_path = str(args.prefetched_input_path or "").strip()
    expected_start = _isoformat_utc(window_start)
    expected_end = _isoformat_utc(window_end)
    if prefetched_input_path:
        prefetched_relative_path, prefetched_prompt = _prepare_prefetched_case_input(
            prefetched_input_path=prefetched_input_path,
            artifact_filename="prefetched_read_iops_input.json",
            expected_start=expected_start,
            expected_end=expected_end,
        )
        evidence_mode = "prefetched_input"
    return {
        "id": f"bjt-{local_date}-{args.local_hour:02d}{args.local_minute:02d}-{args.duration_minutes}m",
        "cluster_id": args.cluster_id,
        "start": expected_start,
        "end": expected_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",
        "evidence_mode": evidence_mode,
        "prefetched_evidence_path": prefetched_relative_path,
        "prefetched_infra_prompt": prefetched_prompt,
    }


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

    namespace = f"tidb{args.parent_id}"
    case_input = _load_runtime_output_data(args.case_input_output_id)
    evidence_mode = str(case_input.get("evidence_mode") or "live_query").strip() or "live_query"
    if evidence_mode == "prefetched_input":
        relative_path = str(case_input.get("prefetched_evidence_path") or "").strip()
        if not relative_path:
            raise MetricsQueryError("case_input is missing prefetched_evidence_path")
        prefetched_input = _load_json_object_from_path(_workspace_path_from_relative(relative_path))
        ranked_nodes, node_to_tikv_node = _prefetched_ranked_nodes(prefetched_input)
        inventory_query = "prefetched_input.inventory"
        per_instance_query = "prefetched_input.per_instance_read_iops"
        metrics_scope: dict[str, Any] = {
            "source": "prefetched_workspace_artifact",
            "path": relative_path,
        }
    else:
        client = O11yMetricsClient(
            provider=args.provider,
            region=args.region,
            infra_uuid=args.infra_uuid,
            github_pat=_env_github_token(),
        )
        midpoint_ts = start_ts + ((end_ts - start_ts) // 2)
        node_to_tikv_node, inventory_query = _resolve_inventory_mapping(
            client,
            namespace=namespace,
            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(
            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,
        )
        metrics_scope = {
            "provider": args.provider,
            "region": args.region,
            "infra_uuid": args.infra_uuid,
            "namespace": namespace,
        }

    return {
        "cluster_id": args.cluster_id,
        "window": {
            "start": _isoformat_utc(start_ts),
            "end": _isoformat_utc(end_ts),
        },
        "mapping_strategy": "node_filter",
        "evidence_sufficient": len(ranked_nodes) >= args.top_n,
        "node_inventory_count": len(node_to_tikv_node),
        "top_n": args.top_n,
        "top_nodes_by_avg_read_iops": ranked_nodes[: args.top_n] if len(ranked_nodes) >= args.top_n else [],
        "queries": {
            "inventory_query": inventory_query,
            "per_instance_query": per_instance_query,
            "metrics_scope": metrics_scope,
        },
        "node_to_tikv_node": node_to_tikv_node,
    }


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Runtime hook for premium TiKV read IOPS ranking")
    subparsers = parser.add_subparsers(dest="command", required=True)

    prepare_parser = subparsers.add_parser("prepare-case-input")
    prepare_parser.add_argument("--cluster-id", required=True)
    prepare_parser.add_argument("--local-hour", type=int, default=11)
    prepare_parser.add_argument("--local-minute", type=int, default=0)
    prepare_parser.add_argument("--duration-minutes", type=int, default=5)
    prepare_parser.add_argument("--step-seconds", type=int, default=300)
    prepare_parser.add_argument("--top-n", type=int, default=3)
    prepare_parser.add_argument("--prefetched-input-path", default="")

    reference_parser = subparsers.add_parser("build-reference")
    reference_parser.add_argument("--cluster-id", required=True)
    reference_parser.add_argument("--parent-id", required=True)
    reference_parser.add_argument("--provider", required=True)
    reference_parser.add_argument("--region", required=True)
    reference_parser.add_argument("--infra-uuid", 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)
    reference_parser.add_argument("--case-input-output-id", default="case_input")

    return parser


def main(argv: list[str]) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)
    if args.command == "prepare-case-input":
        payload = _prepare_case_input(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:]))