"""
Verify that pre-aggregated queries produce the same RU results as raw queries.

Compares: sum(increase(raw_metric{name!="default"}[6h])) by(name, type)
vs:       increase(sum by(name, type)(raw_metric{name!="default"})[6h])

If pre-aggregation is valid, both should produce identical (or near-identical) RU values.
"""

import time
import requests

TEST_URL = "https://www.ds.us-east-1.aws.observability.tidbcloud.com/internal/metrics/f502cc26-7a0e-4a42-a92e-8816c323ccc7/api/v1/"

# 6h window ending recently
END = int(time.time()) - 3600  # 1h ago to ensure data exists
START = END - 21600  # 6h window
STEP = 3600


def query_range(query):
    url = TEST_URL.rstrip("/") + "/query_range"
    params = {"query": query, "start": START, "end": END, "step": STEP}
    resp = requests.post(url, data=params, timeout=300)
    resp.raise_for_status()
    data = resp.json()
    if data.get("status") != "success":
        return {}
    results = data["data"]["result"]
    return results


def extract_series_map(results):
    """Build {label_key: [(ts, value), ...]} from query_range results."""
    out = {}
    for series in results:
        metric = series.get("metric", {})
        name = metric.get("name", "?")
        typ = metric.get("type", "")
        key = f"{name}|{typ}"
        values = [(int(ts), float(v)) for ts, v in series.get("values", [])]
        out[key] = values
    return out


def compute_increase(values):
    """Compute increase (sum of positive deltas) from a time series."""
    total = 0.0
    for i in range(1, len(values)):
        delta = values[i][1] - values[i - 1][1]
        if delta > 0:
            total += delta
    return total


def compare_metric(label, metric_name, agg_labels):
    print(f"\n{'=' * 80}")
    print(f"Metric: {label} ({metric_name})")
    print(f"{'=' * 80}")

    # Raw: get all series, then we'll compute increase per (name, type) group
    raw_results = query_range(f'{metric_name}{{name!="default"}}')
    raw_map = extract_series_map(raw_results)

    # Aggregated: sum by the labels we want
    agg_query = f'sum by ({agg_labels}) ({metric_name}{{name!="default"}})'
    agg_results = query_range(agg_query)
    agg_map = extract_series_map(agg_results)

    print(f"Raw series: {len(raw_results)}, Aggregated series: {len(agg_results)}")

    # For raw: multiple series may share the same (name, type) key
    # Group and sum their increases
    raw_increases = {}
    for key, values in raw_map.items():
        inc = compute_increase(values)
        raw_increases[key] = raw_increases.get(key, 0) + inc

    # For aggregated: one series per (name, type), compute increase directly
    agg_increases = {}
    for key, values in agg_map.items():
        agg_increases[key] = compute_increase(values)

    # Compare top keyspaces
    all_keys = set(raw_increases.keys()) | set(agg_increases.keys())
    diffs = []
    for key in all_keys:
        raw_val = raw_increases.get(key, 0)
        agg_val = agg_increases.get(key, 0)
        if raw_val > 0 or agg_val > 0:
            pct_diff = abs(raw_val - agg_val) / max(raw_val, agg_val, 1) * 100
            diffs.append((key, raw_val, agg_val, pct_diff))

    diffs.sort(key=lambda x: -max(x[1], x[2]))

    print(f"\n{'Keyspace|Type':<25} {'Raw increase':>15} {'Agg increase':>15} {'Diff%':>8}")
    print("-" * 70)
    for key, raw_val, agg_val, pct_diff in diffs[:15]:
        marker = " !!!" if pct_diff > 1.0 else ""
        print(f"{key:<25} {raw_val:>15,.0f} {agg_val:>15,.0f} {pct_diff:>7.2f}%{marker}")

    # Summary
    total_raw = sum(v for _, v, _, _ in diffs)
    total_agg = sum(v for _, _, v, _ in diffs)
    exact_matches = sum(1 for _, r, a, _ in diffs if abs(r - a) < 1)
    close_matches = sum(1 for _, _, _, p in diffs if p < 0.1)

    print(f"\nTotal RU  raw: {total_raw:>15,.0f}")
    print(f"Total RU  agg: {total_agg:>15,.0f}")
    print(f"Exact matches: {exact_matches}/{len(diffs)}")
    print(f"<0.1% diff:    {close_matches}/{len(diffs)}")


def check_cluster_mapping():
    """Verify gateway metric still provides keyspace -> cluster_id mapping."""
    print(f"\n{'=' * 80}")
    print("Cluster mapping: gateway_connection_flow_bytes_total")
    print(f"{'=' * 80}")

    results = query_range('sum by (keyspace_id, cluster_id, type) (gateway_connection_flow_bytes_total)')
    print(f"Aggregated series: {len(results)}")

    mappings = set()
    for series in results:
        m = series.get("metric", {})
        ks = m.get("keyspace_id", "?")
        cl = m.get("cluster_id", "?")
        if ks and cl and ks != "?" and cl != "?":
            mappings.add((ks, cl))

    print(f"Unique keyspace->cluster mappings: {len(mappings)}")
    print("\nSample mappings:")
    for ks, cl in sorted(mappings)[:10]:
        print(f"  keyspace_id={ks} -> cluster_id={cl}")


def main():
    compare_metric("RU Read", "resource_manager_resource_unit_read_request_unit_sum", "name, type")
    compare_metric("RU Write", "resource_manager_resource_unit_write_request_unit_sum", "name, type")
    compare_metric("RU SQL", "resource_manager_resource_unit_sql_layer_request_unit_sum", "name")
    check_cluster_mapping()


if __name__ == "__main__":
    main()