#!/usr/bin/env python3
import statistics
from typing import Any

from common import (
    build_common_arg_parser,
    clinic_cluster_meta,
    clinic_metrics_query,
    clinic_session,
    find_metric_definition,
    inject_expression_variables,
    parse_clinic_link,
    series_last_values,
    write_json,
)


CPU_THRESHOLD = 0.80
MEMORY_THRESHOLD = 0.80
DISK_USED_THRESHOLD = 0.80
DEFAULT_INSTANCE_REGEX = ".*"


def fetch_expression_result(
    session,
    clinic_context: dict[str, Any],
    dashboard_title: str,
    panel_title: str,
    expression_index: int,
    instance: str = DEFAULT_INSTANCE_REGEX,
    k8s_cluster: str = "",
    step: str = "60s",
) -> dict[str, Any]:
    metric_def = find_metric_definition(dashboard_title, panel_title)
    expressions = metric_def["expressions"]
    if expression_index < 0 or expression_index >= len(expressions):
        raise ValueError(
            f"expression_index out of range for dashboard={dashboard_title!r} panel={panel_title!r}: {expression_index}"
        )
    raw_expr = expressions[expression_index]
    final_expr = inject_expression_variables(
        raw_expr,
        clinic_context,
        variable_rules={
            "k8s_cluster": k8s_cluster,
            "instance": instance,
        },
    )
    payload = clinic_metrics_query(session, clinic_context, final_expr, step=step)
    return {
        "knowledge_file": metric_def["knowledge_file"],
        "dashboard_title": metric_def["dashboard_title"],
        "panel_title": metric_def["panel_title"],
        "raw_expression": raw_expr,
        "final_promql": final_expr,
        "raw_series": payload,
        "last_values": series_last_values(payload),
    }


def to_float(value: Any) -> float | None:
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def instance_value_map(last_values: list[dict[str, Any]]) -> dict[str, float]:
    result: dict[str, float] = {}
    for item in last_values:
        instance = item.get("instance")
        value = to_float(item.get("last_value"))
        if not instance or value is None:
            continue
        result[instance] = value
    return result


def scalar_value(last_values: list[dict[str, Any]]) -> float | None:
    for item in last_values:
        value = to_float(item.get("last_value"))
        if value is not None:
            return value
    return None


def ratio_summary(numerator: dict[str, float], denominator: dict[str, float]) -> dict[str, Any]:
    pairs = []
    for instance, num in numerator.items():
        denom = denominator.get(instance)
        if denom in (None, 0):
            continue
        pairs.append((instance, num / denom))
    pairs.sort(key=lambda x: x[1], reverse=True)
    if not pairs:
        return {
            "series_count": 0,
            "max_ratio": None,
            "min_ratio": None,
            "avg_ratio": None,
            "top5": [],
        }
    values = [value for _, value in pairs]
    return {
        "series_count": len(pairs),
        "max_ratio": max(values),
        "min_ratio": min(values),
        "avg_ratio": statistics.mean(values),
        "top5": pairs[:5],
    }


def nonzero_series(last_values: list[dict[str, Any]]) -> list[dict[str, Any]]:
    result = []
    for item in last_values:
        value = to_float(item.get("last_value"))
        if value is None or value <= 0:
            continue
        result.append(
            {
                "instance": item.get("instance"),
                "metric": item.get("metric") or {},
                "last_value": value,
            }
        )
    return result


def check_cpu(session, clinic_context: dict[str, Any]) -> dict[str, Any]:
    usage = fetch_expression_result(session, clinic_context, "Test-Cluster-TiKV-Details", "CPU", 0)
    quota = fetch_expression_result(session, clinic_context, "Test-Cluster-TiKV-Details", "CPU", 1)
    summary = ratio_summary(instance_value_map(usage["last_values"]), instance_value_map(quota["last_values"]))
    max_ratio = summary["max_ratio"]
    passed = bool(max_ratio is not None and max_ratio < CPU_THRESHOLD)
    return {
        "id": "tikv_cpu_usage",
        "title": "TiKV -> Cluster -> CPU",
        "queries": [
            {"name": "cpu_usage", "promql": usage["final_promql"]},
            {"name": "cpu_quota", "promql": quota["final_promql"]},
        ],
        "observed": summary,
        "passed": passed,
        "note": f"Pass if max per-instance CPU usage ratio stays below {CPU_THRESHOLD:.0%}.",
    }


def check_memory(session, clinic_context: dict[str, Any]) -> dict[str, Any]:
    usage = fetch_expression_result(session, clinic_context, "Test-Cluster-TiKV-Details", "Memory", 0)
    quota = fetch_expression_result(session, clinic_context, "Test-Cluster-TiKV-Details", "Memory", 1)
    summary = ratio_summary(instance_value_map(usage["last_values"]), instance_value_map(quota["last_values"]))
    max_ratio = summary["max_ratio"]
    passed = bool(max_ratio is not None and max_ratio < MEMORY_THRESHOLD)
    return {
        "id": "tikv_memory_usage",
        "title": "TiKV -> Cluster -> Memory",
        "queries": [
            {"name": "memory_usage", "promql": usage["final_promql"]},
            {"name": "memory_quota", "promql": quota["final_promql"]},
        ],
        "observed": summary,
        "passed": passed,
        "note": f"Pass if max per-instance memory usage ratio stays below {MEMORY_THRESHOLD:.0%}.",
    }


def check_disk_usage(session, clinic_context: dict[str, Any]) -> dict[str, Any]:
    available = fetch_expression_result(session, clinic_context, "Test-Cluster-TiKV-Details", "Available size", 0)
    capacity = fetch_expression_result(session, clinic_context, "Test-Cluster-TiKV-Details", "Capacity size", 0)
    available_map = instance_value_map(available["last_values"])
    capacity_map = instance_value_map(capacity["last_values"])
    used_ratio_map = {}
    for instance, cap in capacity_map.items():
        avail = available_map.get(instance)
        if avail is None or cap == 0:
            continue
        used_ratio_map[instance] = 1 - (avail / cap)
    summary = ratio_summary(used_ratio_map, {k: 1.0 for k in used_ratio_map})
    max_ratio = summary["max_ratio"]
    passed = bool(max_ratio is not None and max_ratio < DISK_USED_THRESHOLD)
    return {
        "id": "tikv_disk_used_ratio",
        "title": "TiKV -> Cluster -> Available size / Capacity size",
        "queries": [
            {"name": "available_size", "promql": available["final_promql"]},
            {"name": "capacity_size", "promql": capacity["final_promql"]},
        ],
        "observed": summary,
        "passed": passed,
        "note": f"Pass if max per-instance disk used ratio stays below {DISK_USED_THRESHOLD:.0%}.",
    }


def check_busy_threads(session, clinic_context: dict[str, Any]) -> dict[str, Any]:
    result = fetch_expression_result(session, clinic_context, "Test-Cluster-TiKV-Details", "Busy Threads (>80%)", 0)
    busy = nonzero_series(result["last_values"])
    return {
        "id": "tikv_busy_threads",
        "title": "TiKV -> Thread CPU -> Busy Threads (>80%)",
        "queries": [{"name": "busy_threads", "promql": result["final_promql"]}],
        "observed": {
            "busy_thread_count": len(busy),
            "top5": busy[:5],
        },
        "passed": len(busy) == 0,
        "note": "Pass if there are no busy TiKV threads over 80%.",
    }


def check_server_report_failures(session, clinic_context: dict[str, Any]) -> dict[str, Any]:
    result = fetch_expression_result(session, clinic_context, "Test-Cluster-TiKV-Details", "Server report failures", 0)
    failures = nonzero_series(result["last_values"])
    return {
        "id": "tikv_server_report_failures",
        "title": "TiKV -> Errors -> Server report failures",
        "queries": [{"name": "server_report_failures", "promql": result["final_promql"]}],
        "observed": {
            "failure_series_count": len(failures),
            "top5": failures[:5],
        },
        "passed": len(failures) == 0,
        "note": "Pass if the failure report rate stays at 0.",
    }


def check_server_is_busy(session, clinic_context: dict[str, Any]) -> dict[str, Any]:
    subqueries = [
        ("scheduler_too_busy", 0),
        ("channel_full", 1),
        ("coprocessor_full", 2),
        ("store_write_block_wait", 4),
        ("raftstore_process_busy", 5),
    ]
    observed = []
    queries = []
    failures = []
    for name, expr_index in subqueries:
        result = fetch_expression_result(session, clinic_context, "Test-Cluster-TiKV-Details", "Server is busy", expr_index)
        active = nonzero_series(result["last_values"])
        queries.append({"name": name, "promql": result["final_promql"]})
        observed.append(
            {
                "name": name,
                "active_series_count": len(active),
                "top5": active[:5],
            }
        )
        if active:
            failures.append(name)
    return {
        "id": "tikv_server_is_busy",
        "title": "TiKV -> Errors -> Server is busy",
        "queries": queries,
        "observed": observed,
        "passed": len(failures) == 0,
        "note": "Pass if all server-is-busy related rates stay at 0.",
    }


def check_pd_store_count(
    session,
    clinic_context: dict[str, Any],
    check_id: str,
    title: str,
    expr_indices: list[int],
    note: str,
) -> dict[str, Any]:
    queries = []
    observed = []
    total = 0.0
    for expr_index in expr_indices:
        result = fetch_expression_result(session, clinic_context, "PD", "Abnormal stores", expr_index)
        value = scalar_value(result["last_values"]) or 0.0
        total += value
        queries.append(
            {
                "name": f"expr_{expr_index}",
                "promql": result["final_promql"],
            }
        )
        observed.append(
            {
                "expression_index": expr_index,
                "value": value,
            }
        )
    return {
        "id": check_id,
        "title": title,
        "queries": queries,
        "observed": {
            "sub_metrics": observed,
            "total": total,
        },
        "passed": total == 0,
        "note": note,
    }


def run_checks(session, clinic_context: dict[str, Any]) -> list[dict[str, Any]]:
    return [
        check_cpu(session, clinic_context),
        check_memory(session, clinic_context),
        check_disk_usage(session, clinic_context),
        check_busy_threads(session, clinic_context),
        check_server_report_failures(session, clinic_context),
        check_server_is_busy(session, clinic_context),
        check_pd_store_count(
            session,
            clinic_context,
            "pd_slow_stores",
            "PD -> Abnormal stores -> Slow stores",
            [6],
            "Pass if PD reports 0 slow stores.",
        ),
        check_pd_store_count(
            session,
            clinic_context,
            "pd_low_space_stores",
            "PD -> Abnormal stores -> Lowspace stores",
            [2],
            "Pass if PD reports 0 low-space stores.",
        ),
        check_pd_store_count(
            session,
            clinic_context,
            "pd_other_abnormal_stores",
            "PD -> Abnormal stores -> Other stores",
            [0, 1, 3, 4, 5],
            "Pass if disconnected / unhealth / down / offline / tombstone store counts all stay at 0.",
        ),
    ]


def main() -> int:
    parser = build_common_arg_parser(
        "Run the Clinic basic healthy check and output each check item with query, observed value, and pass/fail"
    )
    args = parser.parse_args()

    clinic_context = parse_clinic_link(args.link)
    session = clinic_session(base_url=args.base_url, profile_name=args.profile)
    cluster_meta = clinic_cluster_meta(session, clinic_context)
    checks = run_checks(session, clinic_context)

    passed = [item["id"] for item in checks if item["passed"]]
    failed = [item["id"] for item in checks if not item["passed"]]

    write_json(
        {
            "clinic_context": clinic_context,
            "cluster_meta": {
                "id": cluster_meta.get("id"),
                "name": cluster_meta.get("name"),
                "version": cluster_meta.get("version"),
                "deployType": cluster_meta.get("deployType"),
                "deployTypeV2": cluster_meta.get("deployTypeV2"),
                "topology": cluster_meta.get("topology"),
                "status": cluster_meta.get("status"),
            },
            "thresholds": {
                "cpu_usage_ratio_max": CPU_THRESHOLD,
                "memory_usage_ratio_max": MEMORY_THRESHOLD,
                "disk_used_ratio_max": DISK_USED_THRESHOLD,
                "busy_threads_max": 0,
                "failure_rate_max": 0,
                "abnormal_store_count_max": 0,
            },
            "checks": checks,
            "summary": {
                "total": len(checks),
                "passed": len(passed),
                "failed": len(failed),
                "passed_ids": passed,
                "failed_ids": failed,
            },
            "overall_passed": len(failed) == 0,
            "scope_note": (
                "This basic healthy check is a resource and abnormal-store preflight only. "
                "Passing it does not rule out workload skew, latency spikes, hotspot issues, or deeper TiKV bottlenecks."
            ),
        }
    )
    return 0


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