"""
Benchmark: aggregated PromQL queries across step/window combos and all endpoints.

Tests whether cardinality reduction via sum-by allows finer step sizes
while staying within the cron time budget.
"""

import os
import sys
import time
import concurrent.futures
import requests
import yaml

RETRY_ATTEMPTS = 2
RETRY_BACKOFF = 2
BASE_START = 1774155600
MAX_WORKERS = 6  # parallel endpoint fetches


def fetch(base_url, query, start, end, step):
    url = base_url.rstrip("/") + "/query_range"
    params = {"query": query, "start": start, "end": end, "step": step}
    for attempt in range(1, RETRY_ATTEMPTS + 1):
        try:
            resp = requests.post(url, data=params, timeout=300)
            resp.raise_for_status()
            data = resp.json()
            if data.get("status") != "success":
                return 0, 0, 0
            results = data["data"]["result"]
            series = len(results)
            rows = sum(len(s.get("values", [])) for s in results)
            return series, rows, len(resp.content)
        except Exception:
            if attempt == RETRY_ATTEMPTS:
                return 0, 0, 0
            time.sleep(RETRY_BACKOFF ** attempt)


AGG_METRICS = [
    ("resource_manager_resource_unit_read_request_unit_sum",
     'sum by (name, type) (resource_manager_resource_unit_read_request_unit_sum)'),
    ("resource_manager_resource_unit_write_request_unit_sum",
     'sum by (name, type) (resource_manager_resource_unit_write_request_unit_sum)'),
    ("resource_manager_resource_unit_sql_layer_request_unit_sum",
     'sum by (name) (resource_manager_resource_unit_sql_layer_request_unit_sum)'),
    ("tikv_store_size_bytes",
     'sum by (name, type) (tikv_store_size_bytes)'),
    ("gateway_connection_flow_bytes_total",
     'sum by (keyspace_id, cluster_id, type) (gateway_connection_flow_bytes_total)'),
    ("gateway_connection_sla_total",
     'sum by (keyspace_id, cluster_id, type) (gateway_connection_sla_total)'),
    ("serverless_metering_rcu_provisioned_capacity",
     'sum by (name) (serverless_metering_rcu_provisioned_capacity)'),
    ("tikv_worker_load_data_wru_cost_counter",
     'sum by (name) (tikv_worker_load_data_wru_cost_counter)'),
    ("tidb_vector_search_table_nums_with_vector_index",
     'sum by (name) (tidb_vector_search_table_nums_with_vector_index)'),
    ("tidb_vector_search_table_nums_with_vector_data",
     'sum by (name) (tidb_vector_search_table_nums_with_vector_data)'),
]

RAW_METRICS = [
    ("resource_manager_resource_unit_read_request_unit_sum",
     "resource_manager_resource_unit_read_request_unit_sum"),
    ("resource_manager_resource_unit_write_request_unit_sum",
     "resource_manager_resource_unit_write_request_unit_sum"),
    ("resource_manager_resource_unit_sql_layer_request_unit_sum",
     "resource_manager_resource_unit_sql_layer_request_unit_sum"),
    ("tikv_store_size_bytes", "tikv_store_size_bytes"),
    ("gateway_connection_flow_bytes_total", "gateway_connection_flow_bytes_total"),
    ("gateway_connection_sla_total", "gateway_connection_sla_total"),
    ("serverless_metering_rcu_provisioned_capacity",
     "serverless_metering_rcu_provisioned_capacity"),
    ("tikv_worker_load_data_wru_cost_counter", "tikv_worker_load_data_wru_cost_counter"),
    ("tidb_vector_search_table_nums_with_vector_index",
     "tidb_vector_search_table_nums_with_vector_index"),
    ("tidb_vector_search_table_nums_with_vector_data",
     "tidb_vector_search_table_nums_with_vector_data"),
]


def load_endpoints():
    config_path = os.path.join(os.path.dirname(__file__), "..", "sources.yml")
    with open(config_path) as f:
        return yaml.safe_load(f)["endpoints"]


def fetch_one_metric_all_endpoints(endpoints, query, start, end, step):
    """Fetch one metric across all endpoints in parallel."""
    series = rows = size = 0

    def do_fetch(ep):
        return fetch(ep["url"], query, start, end, step)

    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
        results = pool.map(do_fetch, endpoints)
        for s, r, sz in results:
            series += s
            rows += r
            size += sz
    return series, rows, size


def run_config(endpoints, metrics, window, step, label, mode):
    start = BASE_START
    end_ts = start + window
    total_series = total_rows = total_bytes = 0
    t_start = time.time()

    for i, (_name, query) in enumerate(metrics, 1):
        mt0 = time.time()
        s, r, sz = fetch_one_metric_all_endpoints(endpoints, query, start, end_ts - 1, step)
        dt = time.time() - mt0
        total_series += s
        total_rows += r
        total_bytes += sz
        sys.stderr.write(f"  [{mode}] {label} | {i}/{len(metrics)} {_name}: "
                         f"{s:,} series, {r:,} rows in {dt:.1f}s\n")
        sys.stderr.flush()

    fetch_time = time.time() - t_start
    return total_series, total_rows, total_bytes, fetch_time


def fmt_row(label, mode, series, rows, size_bytes, fetch_time, window):
    load_est = 120 if mode == "agg" else 180
    total_est = fetch_time + load_est
    wins_per_budget = max(1, int(25 * 60 / total_est))
    hours_per_window = window / 3600
    hph = hours_per_window * wins_per_budget * 2
    catch = 100 / max(hph, 0.1)
    feasible = "YES" if hph > 2 else ("MAYBE" if hph >= 1 else "NO")
    return (f"{label:<24} {mode:<5} {series:>9,} {rows:>12,} "
            f"{size_bytes / 1e6:>7.1f} {fetch_time:>7.1f}s {hours_per_window:>8.0f}h "
            f"{wins_per_budget:>8} {hph:>5.0f} {catch:>9.1f}h  {feasible}")


def main():
    endpoints = load_endpoints()
    print(f"Endpoints: {len(endpoints)}", flush=True)
    print(f"Metrics: {len(AGG_METRICS)}", flush=True)
    print(f"Parallel workers: {MAX_WORKERS}", flush=True)

    header = (f"\n{'Config':<24} {'Mode':<5} {'Series':>9} {'Rows':>12} {'Size MB':>8} "
              f"{'Fetch':>8} {'h/win':>8} {'Win/25m':>8} {'h/h':>6} {'Catch100h':>10}  Feasible?")
    print(header, flush=True)
    print("-" * 122, flush=True)

    configs = [
        (21600,   300, "6h / 5min"),
        (21600,   900, "6h / 15min"),
        (21600,  1800, "6h / 30min"),
        (21600,  3600, "6h / 1h"),
        (43200,   900, "12h / 15min"),
        (43200,  1800, "12h / 30min"),
        (43200,  3600, "12h / 1h"),
        (86400,  1800, "24h / 30min"),
        (86400,  3600, "24h / 1h"),
    ]

    # Baseline first: 6h/1h raw
    sys.stderr.write("\n=== BASELINE: 6h / 1h raw ===\n")
    sys.stderr.flush()
    b_s, b_r, b_sz, b_t = run_config(endpoints, RAW_METRICS, 21600, 3600,
                                      "6h / 1h (BASELINE)", "raw")
    print(fmt_row("6h / 1h (BASELINE)", "raw", b_s, b_r, b_sz, b_t, 21600), flush=True)
    print("-" * 122, flush=True)

    for window, step, label in configs:
        sys.stderr.write(f"\n=== {label} agg ===\n")
        sys.stderr.flush()
        s, r, sz, t = run_config(endpoints, AGG_METRICS, window, step, label, "agg")
        print(fmt_row(label, "agg", s, r, sz, t, window), flush=True)

    print(f"\nColumns: h/win = hours of data per window, Win/25m = windows fitting in "
          f"25min budget (fetch+load),")
    print(f"         h/h = hours of data ingested per real hour (2 cron runs/h), "
          f"Catch100h = hours to catch up 100h gap")
    print(f"         Feasible = YES if h/h > 2 (catches up), MAYBE if 1-2, NO if < 1")


if __name__ == "__main__":
    main()