#!/usr/bin/env python3
"""SQL query performance: build indexes and verify with EXPLAIN QUERY PLAN.

Demonstrates index design rules (leftmost-prefix, covering indexes) and how to
read a query plan to confirm indexes are actually used.
"""
from __future__ import annotations

import sqlite3


def setup() -> sqlite3.Connection:
    conn = sqlite3.connect(":memory:")
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    conn.execute("PRAGMA synchronous=NORMAL")

    conn.executescript("""
        CREATE TABLE customers (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            country TEXT NOT NULL,
            signup_date TEXT NOT NULL
        );
        CREATE TABLE orders (
            id INTEGER PRIMARY KEY,
            customer_id INTEGER NOT NULL REFERENCES customers(id),
            total REAL NOT NULL,
            status TEXT NOT NULL,
            created_at TEXT NOT NULL
        );
    """)

    # Seed data
    countries = ["US", "UK", "DE", "FR", "JP", "CA", "AU", "BR"]
    statuses = ["pending", "paid", "shipped", "delivered", "cancelled"]
    import random
    rng = random.Random(1)
    for i in range(1, 20_001):
        conn.execute(
            "INSERT INTO customers (id, name, country, signup_date) VALUES (?,?,?,?)",
            (i, f"customer_{i}", rng.choice(countries), "2024-01-01"),
        )
        conn.execute(
            "INSERT INTO orders (customer_id, total, status, created_at) VALUES (?,?,?,?)",
            (i, round(rng.uniform(5, 500), 2), rng.choice(statuses), "2024-05-01"),
        )
    conn.commit()
    return conn


def plan(conn: sqlite3.Connection, sql: str) -> str:
    rows = conn.execute(f"EXPLAIN QUERY PLAN {sql}").fetchall()
    return "\n".join(f"  {r[0]} {r[1]} {r[2]} {r[3]}" for r in rows)


def main() -> None:
    conn = setup()

    # A realistic query: paid orders from US customers in a date range
    query = """
        SELECT c.name, o.total
        FROM orders o
        JOIN customers c ON c.id = o.customer_id
        WHERE c.country = 'US'
          AND o.status = 'paid'
          AND o.created_at > '2024-05-15'
        ORDER BY o.total DESC
        LIMIT 20
    """

    print("=== Before indexing ===")
    print(plan(conn, query))

    # Correct index design: equality columns first, then range/sort
    conn.execute("CREATE INDEX idx_orders_status ON orders(status)")
    conn.execute("CREATE INDEX idx_orders_created ON orders(created_at)")
    conn.execute("CREATE INDEX idx_customers_country ON customers(country)")
    conn.execute(
        "CREATE INDEX idx_orders_covering ON orders(status, created_at, total, customer_id)"
    )
    conn.execute("ANALYZE")

    print("\n=== After indexing ===")
    print(plan(conn, query))

    print("\n=== Index design notes ===")
    print("- Composite index (a, b): leftmost-prefix rule applies")
    print("  -> serves WHERE a=? AND WHERE a=? AND b=?, NOT WHERE b=?")
    print("- Put equality columns first, then range/sort columns")
    print("- Covering index includes selected columns to avoid table lookups")
    print("- A 'SCAN' on a big table usually means a missing index")


if __name__ == "__main__":
    main()