"""Create/migrate/seed the sample database without replacing existing records."""
import argparse
from pathlib import Path

import psycopg
from psycopg import sql
from psycopg.conninfo import conninfo_to_dict, make_conninfo

from database import connect, database_url
from realistic_seed import populate, reset_sample

DB_FILES = Path(__file__).parent / "db"


def initialize():
    settings = conninfo_to_dict(database_url())
    name = settings.get("dbname", "stat_dashboard_sample")
    try:
        with connect():
            pass
    except psycopg.OperationalError as error:
        if error.sqlstate != '3D000' and 'does not exist' not in str(error):
            raise
        admin = make_conninfo(database_url(), dbname="postgres")
        with psycopg.connect(admin, autocommit=True) as conn:
            conn.execute(sql.SQL("CREATE DATABASE {} ENCODING 'UTF8'").format(sql.Identifier(name)))
    with connect() as conn:
        installed = conn.execute("SELECT to_regclass('analytics.schema_migrations') AS installed").fetchone()["installed"]
        versions = {r["version"] for r in conn.execute("SELECT version FROM analytics.schema_migrations")} if installed else set()
        for migration in sorted(DB_FILES.glob("[0-9]*.sql")):
            if migration.stem not in versions:
                conn.execute(migration.read_text(encoding="utf-8"))
                print(f"Applied migration {migration.stem}.")
    seed()


def seed():
    with connect() as conn:
        if conn.execute("SELECT 1 FROM games WHERE sample_key='willow-v1'").fetchone():
            print("Sample dataset already exists; keeping existing records.")
        else:
            conn.execute("SET LOCAL statement_timeout=0")
            populate(conn)
            print("Seeded three games with 90 days of sample data.")
        conn.execute((DB_FILES/'seed_interactions.sql').read_text(encoding='utf-8'))
    with connect() as conn:
        conn.execute("ANALYZE")
    status()


def reseed():
    with connect() as conn:
        conn.execute('SET LOCAL statement_timeout=0')
        reset_sample(conn)
        populate(conn)
        conn.execute((DB_FILES/'seed_interactions.sql').read_text(encoding='utf-8'))
    with connect() as conn:
        conn.execute('ANALYZE')
    status()


def status():
    with connect() as conn:
        counts = {}
        for table in ("games", "players", "sessions", "player_snapshots", "events", "game_metrics", "activity_runs", "products", "product_sales", "gamepasses", "gamepass_ownership", "interaction_events"):
            counts[table] = conn.execute(sql.SQL("SELECT count(*) AS count FROM {}").format(sql.Identifier(table))).fetchone()["count"]
        print("PostgreSQL sample records: " + ", ".join(f"{k}={v:,}" for k, v in counts.items()))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("command", choices=["init", "seed", "status", "reset-sample"], default="init", nargs="?")
    globals()[{"init": "initialize", "seed": "seed", "status": "status", "reset-sample": "reseed"}[parser.parse_args().command]]()