#!/usr/bin/env python3
"""Deterministic starter benchmark for Chapter 15.

The script intentionally uses only the Python standard library so that readers can run it
before installing heavier engines. It compares a baseline SQLite table with an indexed
variant for a common data-engineering aggregation query.
"""

from __future__ import annotations

import argparse
import csv
import random
import sqlite3
import statistics
import time
from datetime import datetime, timedelta
from pathlib import Path

SEED = 42
REGIONS = ["Tashkent", "Samarkand", "Bukhara", "Fergana", "Andijan", "Nukus"]
EVENT_TYPES = ["view", "cart", "purchase", "refund"]


def generate_events(path: Path, rows: int) -> None:
    random.seed(SEED)
    start = datetime(2026, 1, 1)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["event_id", "event_ts", "customer_id", "region", "event_type", "amount"])
        for event_id in range(1, rows + 1):
            ts = start + timedelta(seconds=random.randint(0, 45 * 24 * 3600))
            event_type = random.choices(EVENT_TYPES, weights=[70, 18, 10, 2], k=1)[0]
            amount = 0.0 if event_type != "purchase" else round(random.uniform(5, 350), 2)
            writer.writerow([
                event_id,
                ts.isoformat(timespec="seconds"),
                random.randint(1, max(1000, rows // 20)),
                random.choice(REGIONS),
                event_type,
                amount,
            ])


def load_sqlite(db_path: Path, csv_path: Path, with_index: bool) -> None:
    if db_path.exists():
        db_path.unlink()
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()
    cur.execute(
        """
        CREATE TABLE events (
            event_id INTEGER PRIMARY KEY,
            event_ts TEXT NOT NULL,
            customer_id INTEGER NOT NULL,
            region TEXT NOT NULL,
            event_type TEXT NOT NULL,
            amount REAL NOT NULL
        )
        """
    )
    with csv_path.open() as f:
        reader = csv.DictReader(f)
        cur.executemany(
            "INSERT INTO events VALUES (:event_id, :event_ts, :customer_id, :region, :event_type, :amount)",
            reader,
        )
    if with_index:
        cur.execute("CREATE INDEX idx_events_type_region_ts ON events(event_type, region, event_ts)")
    conn.commit()
    conn.close()


def benchmark(db_path: Path, trials: int) -> list[float]:
    query = """
    SELECT region,
           substr(event_ts, 1, 10) AS event_date,
           count(*) AS purchase_count,
           round(sum(amount), 2) AS revenue
    FROM events
    WHERE event_type = 'purchase'
    GROUP BY region, event_date
    ORDER BY event_date, region
    """
    timings = []
    for _ in range(trials):
        conn = sqlite3.connect(db_path)
        start = time.perf_counter()
        rows = list(conn.execute(query))
        elapsed = time.perf_counter() - start
        conn.close()
        if not rows:
            raise RuntimeError("benchmark query returned no rows")
        timings.append(elapsed)
    return timings


def summarize(label: str, rows: int, db_path: Path, timings: list[float]) -> dict[str, str]:
    median = statistics.median(timings)
    p95 = sorted(timings)[max(0, int(len(timings) * 0.95) - 1)]
    return {
        "variant": label,
        "input_rows": str(rows),
        "database_mb": f"{db_path.stat().st_size / 1024 / 1024:.2f}",
        "median_seconds": f"{median:.6f}",
        "p95_seconds": f"{p95:.6f}",
        "rows_per_second_at_median": f"{rows / median:.0f}",
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--rows", type=int, default=100_000)
    parser.add_argument("--trials", type=int, default=7)
    args = parser.parse_args()

    out = Path(__file__).parent / "output"
    csv_path = out / "events.csv"
    baseline_db = out / "events_baseline.sqlite"
    indexed_db = out / "events_indexed.sqlite"
    results_path = out / "benchmark_results.csv"

    generate_events(csv_path, args.rows)
    load_sqlite(baseline_db, csv_path, with_index=False)
    load_sqlite(indexed_db, csv_path, with_index=True)

    rows = [
        summarize("baseline", args.rows, baseline_db, benchmark(baseline_db, args.trials)),
        summarize("indexed", args.rows, indexed_db, benchmark(indexed_db, args.trials)),
    ]

    with results_path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)

    print(f"Wrote {results_path}")
    for row in rows:
        print(row)


if __name__ == "__main__":
    main()