#!/usr/bin/env python3
"""Collect query-tuning evidence for a specific SQL digest.

Usage:
    python3 collect_query_evidence.py <cluster_id> <start_ts> <end_ts> [--digest <digest>]

If --digest is not given, automatically picks the TopSQL TiKV CPU-hottest digest.

Outputs structured evidence consumed by tidb-query-tuning SKILL.md (sections Q1-Q5).

Zero external dependencies — uses only stdlib + clinic_api.py from clinic-api skill.
"""
from __future__ import annotations

import math
import os
import sys
from datetime import datetime, timezone

# Locate clinic_api.py from the current workspace skill layout.
_CLINIC_API_SEARCH = [
    os.path.join(os.getcwd(), ".opencode", "skills", "builtin", "clinic-cli", "scripts"),
    os.path.join(os.getcwd(), ".opencode", ".all-skills", "nutshell-skills", "platform", "clinic-api", "scripts"),
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "builtin", "clinic-cli", "scripts"),
    os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "..", "..", "nutshell-skills", "platform", "clinic-api", "scripts",
    ),
    os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "..", "..", "..", "nutshell-skills", "platform", "clinic-api", "scripts",
    ),
]
for _p in _CLINIC_API_SEARCH:
    if os.path.isfile(os.path.join(_p, "clinic_api.py")):
        sys.path.insert(0, _p)
        break

from clinic_api import ClinicAPI  # noqa: E402


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _safe_int(v) -> int:
    try:
        return int(float(str(v)))
    except (ValueError, TypeError):
        return 0


def _safe_float(v) -> float:
    try:
        return float(str(v))
    except (ValueError, TypeError):
        return 0.0


def _prom_results(resp: dict) -> list[dict]:
    return resp.get("data", {}).get("result", [])


def _prom_peak(result: dict) -> float:
    vals = result.get("values", [])
    nums = []
    for v in vals:
        try:
            f = float(v[1])
            if not (math.isnan(f) or math.isinf(f)):
                nums.append(f)
        except (ValueError, IndexError):
            pass
    return max(nums) if nums else 0.0


def _prom_mean(result: dict) -> float:
    vals = result.get("values", [])
    nums = []
    for v in vals:
        try:
            f = float(v[1])
            if not (math.isnan(f) or math.isinf(f)):
                nums.append(f)
        except (ValueError, IndexError):
            pass
    return sum(nums) / len(nums) if nums else 0.0


def _instance_name(result: dict) -> str:
    m = result.get("metric", {})
    return m.get("instance", m.get("component", "unknown"))


def _safe_sql(api: ClinicAPI, cluster_id: str, sql: str, label: str,
              errors: list, timeout: int = 60) -> dict:
    resp = api.query_sql(cluster_id, sql, timeout=timeout)
    if "error" in resp:
        errors.append(f"{label}: {resp['error']}")
        return {}
    return resp


def _safe_range(api: ClinicAPI, cluster_id: str, query: str, start: str,
                end: str, label: str, errors: list) -> list[dict]:
    resp = api.query_range(cluster_id, query, start, end)
    if "error" in resp:
        errors.append(f"{label}: {resp['error']}")
        return []
    return _prom_results(resp)


# ---------------------------------------------------------------------------
# Data collection
# ---------------------------------------------------------------------------

def find_target_digest(api: ClinicAPI, cluster_id: str, start: str, end: str,
                       errors: list) -> tuple[str, str, str]:
    """Find the CPU-hottest SQL digest from TopSQL TiKV.

    Returns (full_digest, source_description, normalized_sql).
    """
    topsql_date = datetime.fromtimestamp(float(start), tz=timezone.utc).strftime("%Y-%m-%d")
    sql = (
        f"SELECT sql_digest, any_value(normalized_sql) AS sample_sql, "
        f"SUM(topsql_cpu_time_ms) AS total_cpu_ms "
        f"FROM clinic_data_proxy.topsql "
        f"WHERE date = '{topsql_date}' AND component = 'tikv' "
        f"AND timestamps BETWEEN {start} AND {end} "
        f"GROUP BY sql_digest ORDER BY total_cpu_ms DESC LIMIT 1"
    )
    resp = _safe_sql(api, cluster_id, sql, "find_digest:topsql_tikv", errors, timeout=120)
    cols = resp.get("columns", [])
    rows = resp.get("rows", [])
    if rows and cols:
        top = dict(zip(cols, rows[0]))
        digest = str(top.get("sql_digest", ""))
        sample = str(top.get("sample_sql", ""))
        cpu_ms = _safe_int(top.get("total_cpu_ms", 0))
        return digest, f"topsql_tikv (total_cpu_ms={cpu_ms})", sample
    return "", "", ""


def resolve_slow_digest(api: ClinicAPI, cluster_id: str, start: str, end: str,
                        topsql_digest: str, errors: list) -> str:
    """Resolve a TopSQL digest (uppercase, 64-char) to its slow_query_logs equivalent.

    TopSQL stores digests in uppercase; slow_query_logs in lowercase.
    Uses LIKE prefix match to bridge the gap.
    """
    date_str = datetime.fromtimestamp(float(start), tz=timezone.utc).strftime("%Y%m%d")
    prefix = topsql_digest[:16]
    sql = (
        f"SELECT digest FROM clinic_data_proxy.slow_query_logs "
        f"WHERE date = '{date_str}' AND time >= {start} AND time <= {end} "
        f"AND LOWER(digest) LIKE LOWER('{prefix}%') "
        f"GROUP BY digest LIMIT 1"
    )
    resp = _safe_sql(api, cluster_id, sql, "resolve_digest", errors)
    cols = resp.get("columns", [])
    rows = resp.get("rows", [])
    if rows and cols:
        return str(dict(zip(cols, rows[0])).get("digest", ""))
    return ""


def collect(api: ClinicAPI, cluster_id: str, start: str, end: str,
            target_digest: str | None = None) -> dict:
    """Collect query-tuning evidence for a specific digest."""

    data: dict = {"errors": []}
    errors = data["errors"]

    date_str = datetime.fromtimestamp(float(start), tz=timezone.utc).strftime("%Y%m%d")
    topsql_date = datetime.fromtimestamp(float(start), tz=timezone.utc).strftime("%Y-%m-%d")

    # ── Step 1: Determine target digest ──
    if target_digest:
        data["target_digest"] = target_digest
        data["digest_source"] = "user-specified"
        data["normalized_sql_preview"] = ""
    else:
        digest, source, sample = find_target_digest(api, cluster_id, start, end, errors)
        if not digest:
            data["target_digest"] = ""
            data["digest_source"] = "not found"
            return data
        data["target_digest"] = digest
        data["digest_source"] = source
        data["normalized_sql_preview"] = sample

    topsql_digest = data["target_digest"]

    # ── Step 2: TopSQL CPU attribution for this digest (per-instance) ──
    topsql_instance_sql = (
        f"SELECT instance, SUM(topsql_cpu_time_ms) AS total_cpu_ms "
        f"FROM clinic_data_proxy.topsql "
        f"WHERE date = '{topsql_date}' AND component = 'tikv' "
        f"AND timestamps BETWEEN {start} AND {end} "
        f"AND sql_digest = '{topsql_digest}' "
        f"GROUP BY instance ORDER BY total_cpu_ms DESC LIMIT 10"
    )
    data["topsql_by_instance"] = _safe_sql(
        api, cluster_id, topsql_instance_sql, "Q1:topsql_by_instance", errors, timeout=120)

    # Also get total CPU per instance (all digests) for percentage calculation
    topsql_total_sql = (
        f"SELECT instance, SUM(topsql_cpu_time_ms) AS total_cpu_ms "
        f"FROM clinic_data_proxy.topsql "
        f"WHERE date = '{topsql_date}' AND component = 'tikv' "
        f"AND timestamps BETWEEN {start} AND {end} "
        f"GROUP BY instance ORDER BY total_cpu_ms DESC LIMIT 50"
    )
    data["topsql_instance_totals"] = _safe_sql(
        api, cluster_id, topsql_total_sql, "Q1:topsql_totals", errors, timeout=120)

    # ── Step 3: Resolve slow_query digest and get full details ──
    slow_digest = resolve_slow_digest(api, cluster_id, start, end, topsql_digest, errors)
    data["slow_digest"] = slow_digest

    if slow_digest:
        # Q2: Full SQL + decoded_plan + stats + backoff for worst slow query
        detail_sql = (
            f"SELECT any_value(query) AS full_sql, "
            f"any_value(decoded_plan) AS decoded_plan, "
            f"any_value(stats) AS stats_info, "
            f"any_value(backoff_types) AS backoff_types, "
            f"any_value(index_names) AS index_names, "
            f"MAX(query_time) AS max_query_time, "
            f"MAX(process_time) AS max_process_time, "
            f"MAX(wait_time) AS max_wait_time, "
            f"MAX(total_keys) AS max_total_keys, "
            f"MAX(process_keys) AS max_process_keys, "
            f"MAX(mem_max) AS max_mem, "
            f"MAX(result_rows) AS max_result_rows "
            f"FROM clinic_data_proxy.slow_query_logs "
            f"WHERE date = '{date_str}' AND time >= {start} AND time <= {end} "
            f"AND digest = '{slow_digest}' "
            f"LIMIT 1"
        )
        data["query_detail"] = _safe_sql(api, cluster_id, detail_sql, "Q2:query_detail", errors)

        # Q3: Plan stability — top 5 slow queries for this digest
        stability_sql = (
            f"SELECT query_time, process_time, wait_time, "
            f"total_keys, process_keys, result_rows, mem_max, "
            f"backoff_types, index_names, instance "
            f"FROM clinic_data_proxy.slow_query_logs "
            f"WHERE date = '{date_str}' AND time >= {start} AND time <= {end} "
            f"AND digest = '{slow_digest}' "
            f"ORDER BY query_time DESC LIMIT 5"
        )
        data["plan_stability"] = _safe_sql(api, cluster_id, stability_sql, "Q3:plan_stability", errors)
    else:
        # Fallback: get normalized_sql from topsql if slow_query has no match
        norm_sql = (
            f"SELECT any_value(normalized_sql) AS full_sql "
            f"FROM clinic_data_proxy.topsql "
            f"WHERE date = '{topsql_date}' AND component = 'tikv' "
            f"AND timestamps BETWEEN {start} AND {end} "
            f"AND sql_digest = '{topsql_digest}' LIMIT 1"
        )
        data["query_detail_fallback"] = _safe_sql(
            api, cluster_id, norm_sql, "Q2:topsql_fallback", errors)

    # ── Step 4: Hotspot node metrics ──
    # Identify hotspot instances from TopSQL
    hotspot_instances = []
    ti_cols = data.get("topsql_by_instance", {}).get("columns", [])
    ti_rows = data.get("topsql_by_instance", {}).get("rows", [])
    if ti_rows and ti_cols:
        for row in ti_rows[:3]:
            rec = dict(zip(ti_cols, row))
            inst = str(rec.get("instance", ""))
            if inst:
                hotspot_instances.append(inst)

    if hotspot_instances:
        inst_regex = "|".join(hotspot_instances)

        # unified_read_pool CPU on hotspot nodes
        data["read_pool_cpu"] = _safe_range(
            api, cluster_id,
            f'sum(rate(tikv_thread_cpu_seconds_total{{name=~"unified_read.*",instance=~"{inst_regex}"}}[1m])) by (instance)',
            start, end, "Q4:read_pool_cpu", errors)

        # coprocessor P99 on hotspot nodes
        data["cop_p99"] = _safe_range(
            api, cluster_id,
            f'histogram_quantile(0.99, sum(rate(tikv_grpc_msg_duration_seconds_bucket{{type="coprocessor",instance=~"{inst_regex}"}}[1m])) by (le, instance))',
            start, end, "Q4:cop_p99", errors)

        # gRPC P99 on hotspot nodes (for cross-reference)
        data["grpc_p99"] = _safe_range(
            api, cluster_id,
            f'histogram_quantile(0.99, sum(rate(tikv_grpc_msg_duration_seconds_bucket{{type!="raft",instance=~"{inst_regex}"}}[1m])) by (le, instance))',
            start, end, "Q4:grpc_p99", errors)

    return data


# ---------------------------------------------------------------------------
# Summary formatting
# ---------------------------------------------------------------------------

def format_summary(data: dict) -> str:
    lines: list[str] = []

    target = data.get("target_digest", "")
    source = data.get("digest_source", "")
    preview = data.get("normalized_sql_preview", "")

    if not target:
        lines.append("## ERROR: No target digest found")
        if data.get("errors"):
            for e in data["errors"]:
                lines.append(f"  - {e}")
        return "\n".join(lines)

    # ── Q1: Target Query ──
    lines.append("## Q1. Target Query")
    lines.append(f"topsql_digest={target}")
    slow_digest = data.get("slow_digest", "")
    if slow_digest:
        lines.append(f"slow_query_digest={slow_digest}")
    lines.append(f"source={source}")

    # Per-instance CPU breakdown
    ti_cols = data.get("topsql_by_instance", {}).get("columns", [])
    ti_rows = data.get("topsql_by_instance", {}).get("rows", [])
    total_cols = data.get("topsql_instance_totals", {}).get("columns", [])
    total_rows = data.get("topsql_instance_totals", {}).get("rows", [])

    # Build instance total map
    inst_totals: dict[str, int] = {}
    for row in total_rows:
        rec = dict(zip(total_cols, row))
        inst_totals[str(rec.get("instance", ""))] = _safe_int(rec.get("total_cpu_ms", 0))

    if ti_rows:
        lines.append("TopSQL TiKV per-instance:")
        for row in ti_rows:
            rec = dict(zip(ti_cols, row))
            inst = str(rec.get("instance", ""))
            cpu_ms = _safe_int(rec.get("total_cpu_ms", 0))
            total = inst_totals.get(inst, 0)
            pct = round(cpu_ms * 100.0 / total, 1) if total > 0 else 0
            lines.append(f"  {inst}: cpu_ms={cpu_ms} cpu_pct={pct}% (of instance total {total})")

    # ── Q2: Execution Plan Analysis ──
    lines.append("\n## Q2. Execution Plan Analysis")
    detail = data.get("query_detail", {})
    detail_cols = detail.get("columns", [])
    detail_rows = detail.get("rows", [])
    if detail_rows:
        rec = dict(zip(detail_cols, detail_rows[0]))
        full_sql = str(rec.get("full_sql", ""))
        lines.append(f"full_sql: {full_sql}")
        lines.append(f"stats_info: {rec.get('stats_info', '')}")
        lines.append(f"backoff_types: {rec.get('backoff_types', '')}")
        lines.append(f"index_names: {rec.get('index_names', '')}")
        lines.append(
            f"query_time={rec.get('max_query_time', '')} "
            f"process_time={rec.get('max_process_time', '')} "
            f"wait_time={rec.get('max_wait_time', '')} "
            f"total_keys={rec.get('max_total_keys', '')} "
            f"process_keys={rec.get('max_process_keys', '')} "
            f"mem_max={rec.get('max_mem', '')} "
            f"result_rows={rec.get('max_result_rows', '')}"
        )
        decoded = str(rec.get("decoded_plan", ""))
        if decoded:
            lines.append(f"\ndecoded_plan:\n{decoded}")
        else:
            lines.append("\ndecoded_plan: (empty)")
    else:
        # Fallback to topsql normalized_sql
        fb = data.get("query_detail_fallback", {})
        fb_cols = fb.get("columns", [])
        fb_rows = fb.get("rows", [])
        if fb_rows:
            rec = dict(zip(fb_cols, fb_rows[0]))
            lines.append(f"full_sql (from topsql, normalized): {rec.get('full_sql', '')}")
        else:
            lines.append(f"normalized_sql_preview: {preview}")
        lines.append("(no slow_query match — decoded_plan unavailable)")

    # ── Q3: Plan Stability ──
    lines.append("\n## Q3. Plan Stability (top 5 slow queries for this digest)")
    stab = data.get("plan_stability", {})
    stab_cols = stab.get("columns", [])
    stab_rows = stab.get("rows", [])
    if stab_rows:
        for i, row in enumerate(stab_rows):
            rec = dict(zip(stab_cols, row))
            lines.append(
                f"  [{i+1}] query_time={rec.get('query_time', '')} "
                f"process_time={rec.get('process_time', '')} "
                f"wait_time={rec.get('wait_time', '')} "
                f"total_keys={rec.get('total_keys', '')} "
                f"process_keys={rec.get('process_keys', '')} "
                f"result_rows={rec.get('result_rows', '')} "
                f"mem_max={rec.get('mem_max', '')} "
                f"backoff={rec.get('backoff_types', '')} "
                f"instance={rec.get('instance', '')}"
            )
    else:
        lines.append("  (no slow queries found for this digest in window)")

    # ── Q4: Hotspot Node Metrics ──
    lines.append("\n## Q4. Hotspot Node Metrics")

    read_pool = data.get("read_pool_cpu", [])
    if read_pool:
        lines.append("unified_read_pool CPU:")
        for r in sorted(read_pool, key=_prom_peak, reverse=True):
            lines.append(f"  {_instance_name(r)}: peak={_prom_peak(r):.2f} cores mean={_prom_mean(r):.2f}")

    cop = data.get("cop_p99", [])
    if cop:
        lines.append("coprocessor P99:")
        for r in sorted(cop, key=_prom_peak, reverse=True):
            lines.append(f"  {_instance_name(r)}: peak={_prom_peak(r):.3f}s mean={_prom_mean(r):.3f}s")

    grpc = data.get("grpc_p99", [])
    if grpc:
        lines.append("gRPC P99 (non-raft):")
        for r in sorted(grpc, key=_prom_peak, reverse=True):
            lines.append(f"  {_instance_name(r)}: peak={_prom_peak(r):.3f}s mean={_prom_mean(r):.3f}s")

    if not read_pool and not cop and not grpc:
        lines.append("  (no hotspot instances identified)")

    # ── Errors ──
    if data.get("errors"):
        lines.append("\n## Collection Errors (non-fatal)")
        for e in data["errors"]:
            lines.append(f"  - {e}")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    if len(sys.argv) < 4:
        print("Usage: python3 collect_query_evidence.py <cluster_id> <start_ts> <end_ts> [--digest <digest>]")
        print("  start_ts/end_ts: Unix timestamps (seconds)")
        print("  --digest: specific SQL digest to analyze (optional; auto-detects TopSQL TiKV CPU-hottest)")
        print()
        print("Example:")
        print("  python3 collect_query_evidence.py 10754652599988043412 1743825600 1743832800")
        print("  python3 collect_query_evidence.py 10754652599988043412 1743825600 1743832800 --digest 61B6924329F86D8B...")
        sys.exit(1)

    cluster_id = sys.argv[1]
    start = sys.argv[2]
    end = sys.argv[3]

    target_digest = None
    if "--digest" in sys.argv:
        idx = sys.argv.index("--digest")
        if idx + 1 < len(sys.argv):
            target_digest = sys.argv[idx + 1]

    api = ClinicAPI()

    print(f"Collecting query-tuning evidence for cluster {cluster_id} ({start} to {end})...",
          file=sys.stderr)
    if target_digest:
        print(f"Target digest: {target_digest}", file=sys.stderr)
    else:
        print("Auto-detecting TopSQL TiKV CPU-hottest digest...", file=sys.stderr)

    data = collect(api, cluster_id, start, end, target_digest)
    print(format_summary(data))


if __name__ == "__main__":
    main()