from __future__ import annotations

import argparse
import json
import math
import os
import sys
import urllib.error
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")
TRIVIAL_SQL_PREFIXES = (
    "commit",
    "start transaction",
    "begin",
    "rollback",
    "set",
    "use",
    "show",
)


class SlowQueryReferenceError(RuntimeError):
    pass


class ClinicDataProxyClient:
    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_sql(self, *, cluster_id: str, sql: str, timeout_s: int = DEFAULT_TIMEOUT_S) -> dict[str, Any]:
        url = f"{self._base_url}/data-proxy/query"
        payload = json.dumps(
            {
                "clusterId": cluster_id,
                "timeout": timeout_s,
                "sql": sql,
            }
        ).encode("utf-8")
        request = urllib.request.Request(
            url,
            data=payload,
            headers={
                "Authorization": f"Bearer {self._api_key}",
                "Content-Type": "application/json",
            },
            method="POST",
        )
        try:
            with urllib.request.urlopen(request, timeout=self._timeout_s) as response:
                body = response.read().decode("utf-8")
        except urllib.error.HTTPError as exc:
            text = exc.read().decode("utf-8", errors="replace")
            raise SlowQueryReferenceError(f"HTTP {exc.code} for /data-proxy/query: {text}") from exc
        except urllib.error.URLError as exc:
            raise SlowQueryReferenceError(f"request failed for /data-proxy/query: {exc}") from exc
        try:
            data = json.loads(body)
        except json.JSONDecodeError as exc:
            raise SlowQueryReferenceError(f"invalid JSON response for /data-proxy/query: {exc}") from exc
        if isinstance(data, dict) and data.get("rows") is None:
            data["rows"] = []
        error_text = str(data.get("error") or "").strip() if isinstance(data, dict) else ""
        if error_text:
            raise SlowQueryReferenceError(error_text)
        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 SlowQueryReferenceError(
        "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 _safe_int(value: Any) -> int | None:
    numeric = _safe_float(value)
    if numeric is None:
        return None
    return int(numeric)


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 _sql_date_partition(start_ts: int) -> str:
    return datetime.fromtimestamp(start_ts, tz=UTC).strftime("%Y%m%d")


def _business_sql_normalized(sql_text: str) -> str:
    return " ".join(sql_text.strip().lower().split())


def _is_trivial_control_sql(sql_text: str) -> bool:
    normalized = _business_sql_normalized(sql_text)
    if not normalized:
        return True
    return any(
        normalized == prefix
        or normalized.startswith(f"{prefix} ")
        or normalized.startswith(f"{prefix};")
        for prefix in TRIVIAL_SQL_PREFIXES
    )


def _build_slow_query_sql(*, start_ts: int, end_ts: int, limit: int) -> str:
    return (
        "SELECT digest, any_value(query) AS sample_sql, COUNT(*) AS exec_count, "
        "MAX(query_time) AS max_time, AVG(query_time) AS avg_time, MAX(mem_max) AS max_mem "
        "FROM clinic_data_proxy.slow_query_logs "
        f"WHERE date = '{_sql_date_partition(start_ts)}' AND time >= {start_ts} AND time <= {end_ts} "
        "GROUP BY digest ORDER BY exec_count DESC LIMIT "
        f"{limit}"
    )


def _parse_reference_rows(payload: dict[str, Any], *, top_n: int) -> tuple[list[dict[str, Any]], int]:
    rows = payload.get("rows")
    if not isinstance(rows, list):
        raise SlowQueryReferenceError("slow query SQL response is missing rows")

    candidates: list[dict[str, Any]] = []
    excluded = 0
    for row in rows:
        if not isinstance(row, list) or len(row) < 6:
            continue
        digest = str(row[0] or "").strip().lower()
        sample_sql = str(row[1] or "").strip()
        exec_count = _safe_int(row[2])
        avg_time_s = _safe_float(row[4])
        if not digest or exec_count is None or avg_time_s is None:
            continue
        if _is_trivial_control_sql(sample_sql):
            excluded += 1
            continue
        candidates.append(
            {
                "digest": digest,
                "exec_count": exec_count,
                "avg_time_s": round(avg_time_s, 3),
                "sample_sql_preview": sample_sql[:180],
            }
        )

    candidates.sort(
        key=lambda item: (
            -int(item["exec_count"]),
            -float(item["avg_time_s"]),
            str(item["digest"]),
        )
    )
    if len(candidates) < top_n:
        raise SlowQueryReferenceError(
            f"expected at least {top_n} business slow queries after filtering, got {len(candidates)}"
        )
    return candidates[:top_n], excluded


def _resolve_window(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,
    )
    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(window_start),
        "end": _isoformat_utc(window_end),
        "local_timezone": "Asia/Shanghai",
        "local_date": local_date,
        "duration_minutes": args.duration_minutes,
        "top_n": args.top_n,
        "filter_rule": "exclude_transaction_and_control_sql",
        "ranking_rule": "exec_count_desc_then_avg_time_desc_then_digest_asc",
    }


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

    query_sql = _build_slow_query_sql(start_ts=start_ts, end_ts=end_ts, limit=max(args.top_n * 10, 30))
    client = ClinicDataProxyClient(
        api_key=_env_api_key(),
        base_url=os.environ.get("CLINIC_ENDPOINT", DEFAULT_BASE_URL).strip() or DEFAULT_BASE_URL,
    )
    payload = client.query_sql(cluster_id=args.cluster_id, sql=query_sql)
    slow_queries, excluded = _parse_reference_rows(payload, top_n=args.top_n)

    return {
        "cluster_id": args.cluster_id,
        "cluster_label": args.cluster_label,
        "window": {
            "start": _isoformat_utc(start_ts),
            "end": _isoformat_utc(end_ts),
        },
        "evidence_sufficient": True,
        "slow_queries": [
            {
                "digest": str(item["digest"]),
                "exec_count": int(item["exec_count"]),
                "avg_time_s": float(item["avg_time_s"]),
            }
            for item in slow_queries
        ],
        "reference_rows": slow_queries,
        "debug": {
            "query_sql": query_sql,
            "raw_row_count": len(payload.get("rows") or []),
            "excluded_control_sql_count": excluded,
            "filter_rule": "exclude_transaction_and_control_sql",
            "ranking_rule": "exec_count_desc_then_avg_time_desc_then_digest_asc",
        },
    }


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Runtime hook for representative business slow queries in the premium cluster"
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    resolve_parser = subparsers.add_parser("resolve-window")
    resolve_parser.add_argument("--cluster-id", required=True)
    resolve_parser.add_argument("--cluster-label", required=True)
    resolve_parser.add_argument("--local-hour", type=int, default=11)
    resolve_parser.add_argument("--local-minute", type=int, default=0)
    resolve_parser.add_argument("--duration-minutes", type=int, default=5)
    resolve_parser.add_argument("--top-n", type=int, default=3)

    reference_parser = subparsers.add_parser("build-reference")
    reference_parser.add_argument("--cluster-id", required=True)
    reference_parser.add_argument("--cluster-label", required=True)
    reference_parser.add_argument("--start", required=True)
    reference_parser.add_argument("--end", required=True)
    reference_parser.add_argument("--top-n", type=int, default=3)

    return parser


def main(argv: list[str]) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)
    if args.command == "resolve-window":
        payload = _resolve_window(args)
    elif args.command == "build-reference":
        payload = _build_reference(args)
    else:
        raise SlowQueryReferenceError(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:]))