#!/usr/bin/env python3
"""
Inspect a TiDB Cloud cluster for the last hour: metrics, log-stats, slowlog-stats.
Produces an anomaly report.

Usage:
  python skills/o11y-data/inspect_cluster.py <cluster_id> [--region us-west-2]
"""
import sys
import json
import sqlite3
import tempfile
from datetime import datetime, timezone, timedelta


def main():
    cluster_id = sys.argv[1] if len(sys.argv) > 1 else None
    region = "us-west-2"
    if "--region" in sys.argv:
        i = sys.argv.index("--region")
        region = sys.argv[i + 1]

    if not cluster_id:
        print("Usage: python inspect_cluster.py <cluster_id> [--region us-west-2]")
        sys.exit(1)

    import requests
    import boto3

    from auth.client_auth import get_credentials

    creds = get_credentials()
    auth0_token = creds.get("auth0_token")
    if not auth0_token:
        raise SystemExit("Failed to get auth0_token")

    # 1. Get cluster info (o11y_id)
    # Try given region first, then common AWS regions and bizTypes
    biz_types = ["tidbcloud", "prod-tidbcloud-nextgen", "tidbcloud-nextgen"]
    regions = [region, "us-east-1", "ap-northeast-1", "ap-southeast-1"]
    clusters = []
    for bt in biz_types:
        for reg in regions:
            url = f"http://www.gs.us-west-2.aws.observability.tidbcloud.com/api/v1/clusters/?bizType={bt}&vendor=aws&region={reg}&applicationID={cluster_id}"
            r = requests.get(url, headers={"Authorization": f"Bearer {auth0_token}"})
            if r.status_code != 200:
                continue
            data = r.json()
            lst = data.get("data", {}).get("clusters", [])
            if lst:
                clusters = lst
                region = reg
                break
        if clusters:
            break
    if not clusters:
        raise SystemExit(f"No cluster found for {cluster_id} (tried regions: {regions}, bizTypes: {biz_types})")

    cl = clusters[0]
    o11y_id = cl["id"]
    version = cl.get("version", "?")

    print(f"Cluster: {cluster_id} | o11y_id: {o11y_id} | region: {region} | version: {version}")

    end_ts = int(datetime.now(timezone.utc).timestamp())
    start_ts = end_ts - 3600

    metrics_base = f"https://www.ds.{region}.aws.observability.tidbcloud.com/internal/metrics/{o11y_id}/api/v1/query_range"
    headers = {"Authorization": f"Bearer {auth0_token}"}

    report = {
        "cluster_id": cluster_id,
        "o11y_id": o11y_id,
        "region": region,
        "time_range": f"{datetime.fromtimestamp(start_ts, tz=timezone.utc).isoformat()} to {datetime.fromtimestamp(end_ts, tz=timezone.utc).isoformat()}",
        "anomalies": [],
        "summary": {},
        "details": {},
    }

    # 2. Query metrics
    def query_metrics(promql, name):
        import urllib.parse
        q = urllib.parse.quote(promql)
        u = f"{metrics_base}?query={q}&start={start_ts}&end={end_ts}&step=60"
        r = requests.get(u, headers=headers)
        if r.status_code != 200:
            return None, f"HTTP {r.status_code}"
        j = r.json()
        results = j.get("data", {}).get("result", [])
        return results, None

    # QPS
    results, err = query_metrics("sum(rate(tidb_server_query_total[5m])) by (type)", "qps")
    if results:
        total_qps = 0
        by_type = {}
        for s in results:
            t = s.get("metric", {}).get("type", "unknown")
            vs = s.get("values", [])
            if vs:
                last_val = float(vs[-1][1])
                total_qps += last_val
                by_type[t] = round(last_val, 2)
        report["details"]["qps"] = {"total": round(total_qps, 2), "by_type": by_type}
        if total_qps == 0:
            report["anomalies"].append("QPS is 0 - cluster may be idle or down")
    else:
        report["details"]["qps"] = {"error": err}

    # Latency P99
    results, err = query_metrics(
        "histogram_quantile(0.99, sum(rate(tidb_server_handle_query_duration_seconds_bucket[5m])) by (le))",
        "p99",
    )
    if results:
        vals = [float(v[1]) for s in results for v in s.get("values", []) if float(v[1]) < 1e6]
        p99_sec = max(vals) if vals else 0
        report["details"]["latency_p99_seconds"] = round(p99_sec, 4)
        if p99_sec > 10:
            report["anomalies"].append(f"High P99 latency: {p99_sec:.2f}s")
    else:
        report["details"]["latency"] = {"error": err}

    # Connections
    results, err = query_metrics("sum(tidb_server_connections)", "connections")
    if results:
        vs = [float(v[1]) for s in results for v in s.get("values", [])]
        conn = max(vs) if vs else 0
        report["details"]["connections"] = int(conn)
        if conn > 10000:
            report["anomalies"].append(f"High connection count: {int(conn)}")
    else:
        report["details"]["connections"] = {"error": err}

    # TiKV CPU
    results, err = query_metrics("sum(rate(tikv_thread_cpu_seconds_total[1m])) by (name)", "tikv_cpu")
    if results:
        total = sum(float(v[1]) for s in results for v in s.get("values", [])[-1:])
        report["details"]["tikv_cpu_total"] = round(total, 2)
    else:
        report["details"]["tikv_cpu"] = {"error": err}

    # TiKV regions
    results, err = query_metrics("sum(tikv_raftstore_region_count) by (type)", "regions")
    if results:
        region_counts = {s.get("metric", {}).get("type", "?"): int(float(s.get("values", [[0, 0]])[-1][1])) for s in results}
        report["details"]["tikv_regions"] = region_counts

    # 3. Log-stats from S3
    bucket = f"o11y-prod-shared-{region}"
    hour_utc = datetime.utcnow() - timedelta(hours=1)
    hour_key = hour_utc.strftime("%Y%m%d%H")
    log_stats_key = f"diagnosis/data/{cluster_id}/log-stats/1h/{hour_key}.db"

    s3 = boto3.client(
        "s3",
        aws_access_key_id=creds["access_key_id"],
        aws_secret_access_key=creds["secret_access_key"],
        aws_session_token=creds.get("session_token"),
        region_name=region,
    )

    try:
        with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
            s3.download_file(bucket, log_stats_key, f.name)
            conn = sqlite3.connect(f.name)
            cur = conn.cursor()
            cur.execute(
                "SELECT component, level, SUM(count) as total FROM log_stats GROUP BY component, level ORDER BY total DESC"
            )
            rows = cur.fetchall()
            cur.execute(
                "SELECT component, level, msg_digest, count FROM log_stats WHERE level IN ('ERROR','WARN') ORDER BY count DESC LIMIT 10"
            )
            top_errors = cur.fetchall()
            conn.close()

        summary = {}
        errors = []
        for comp, level, cnt in rows:
            k = f"{comp}_{level}"
            summary[k] = cnt
            if level in ("ERROR", "WARN") and cnt > 100:
                errors.append(f"{comp} {level}: {cnt}")
        report["details"]["log_stats"] = {"hour": hour_key, "summary": summary}
        if errors:
            report["anomalies"].extend(errors)
        report["details"]["top_errors"] = [{"component": r[0], "level": r[1], "msg": (r[2] or "")[:80], "count": r[3]} for r in top_errors]
    except Exception as e:
        report["details"]["log_stats"] = {"error": str(e)}

    # 4. Slowlog-stats (optional, schema may vary)
    slow_key = f"diagnosis/data/{cluster_id}/slowlog-stats/1h/{hour_key}.db"
    try:
        with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
            s3.download_file(bucket, slow_key, f.name)
            conn = sqlite3.connect(f.name)
            cur = conn.cursor()
            cur.execute("SELECT name FROM sqlite_master WHERE type='table'")
            tables = [r[0] for r in cur.fetchall()]
            slow_rows = []
            for t in tables:
                if "slow" in t.lower() or "query" in t.lower():
                    try:
                        cur.execute(f"SELECT * FROM {t} LIMIT 5")
                        slow_rows = cur.fetchall()
                        report["details"]["slowlog_stats"] = {
                            "hour": hour_key,
                            "table": t,
                            "sample_rows": len(slow_rows),
                        }
                        break
                    except sqlite3.OperationalError:
                        pass
            conn.close()
        if "slowlog_stats" not in report["details"]:
            report["details"]["slowlog_stats"] = {"hour": hour_key, "info": "no queryable table"}
    except Exception as e:
        report["details"]["slowlog_stats"] = {"error": str(e)}

    # Summary
    report["summary"]["anomaly_count"] = len(report["anomalies"])
    report["summary"]["status"] = "OK" if not report["anomalies"] else "ANOMALIES"

    # Print report
    print("\n" + "=" * 60)
    print("INSPECTION REPORT")
    print("=" * 60)
    print(f"Time range: {report['time_range']}")
    print(f"Status: {report['summary']['status']} ({report['summary']['anomaly_count']} anomalies)")
    print()
    print("Details:")
    print(json.dumps(report["details"], indent=2, default=str))
    print()
    if report["anomalies"]:
        print("Anomalies:")
        for a in report["anomalies"]:
            print(f"  - {a}")
    print("=" * 60)

    with open(f"inspect_{cluster_id}_{int(end_ts)}.json", "w") as f:
        json.dump(report, f, indent=2, default=str)
    print(f"\nReport saved to inspect_{cluster_id}_{int(end_ts)}.json")


if __name__ == "__main__":
    main()