import statistics as st
from time import sleep, monotonic
from argparse import ArgumentParser, Namespace

import psycopg
import psycopg2
from psycopg import sql

DSN = "dbname=psycopg3_test host=localhost sslmode=disable"

opt: Namespace


def main():
    global opt
    opt = parse_cmdline()

    for i, smodule in enumerate(opt.modules):
        if i == 0 and not opt.no_schema:
            setup_testtable()
        module = globals()[smodule]
        bench_module(module)


def bench_module(module):
    if module is psycopg2:
        print(f"testing {module.__name__} {module.__version__}")
    else:
        if psycopg._cmodule._psycopg:
            vmod = __import__(psycopg._cmodule._psycopg.__name__.split(".")[0])
        else:
            vmod = psycopg
        print(f"testing {module.__name__}[{psycopg.pq.__impl__}] {vmod.__version__}")

    with module.connect(DSN) as conn:
        conn.autocommit = True

        bench_query(conn)


def bench_query(conn):
    query = (
        sql.SQL("select {fields} from testtable limit {limit}")
        .format(
            fields=sql.SQL(",").join(sql.Identifier(f"c{n}") for n in range(opt.ncols)),
            limit=opt.nrows,
        )
        .as_string()
    )
    try:
        psycopg._cmodule._psycopg._load_rows_page_size = opt.page_size
    except AttributeError:
        pass

    with conn.cursor() as cur:
        measure_query(cur, query)


def measure_query(cur, query):
    durations = []
    for _ in range(opt.repeat):
        cur.execute(query)
        start = monotonic()
        cur.fetchall()
        end = monotonic()
        durations.append(end - start)
        if opt.relax:
            sleep(opt.relax)

    mean = st.mean(durations)
    # median = st.median(durations)
    sdev = st.stdev(durations)
    cv = sdev / mean
    msg = f"rows: {opt.nrows} cols: {opt.ncols}"
    try:
        msg += f" page: {cur._tx._load_rows_page_size}"
    except AttributeError:
        pass
    print(
        f"{msg}: {round(mean, 5)} ±{round(sdev, 5)}"
        f" (cv: {round(cv, 3)}, repeats: {opt.repeat})"
    )


def parse_cmdline() -> Namespace:
    parser = ArgumentParser(description=__doc__)
    parser.add_argument(
        "-m",
        "--module",
        nargs="+",
        dest="modules",
        choices=("psycopg2", "psycopg"),
        default=["psycopg2", "psycopg"],
        help="the modules to test [default: %(default)s]",
    )
    parser.add_argument(
        "-r", "--nrows", type=int, default=1000, help="number of rows to select"
    )
    parser.add_argument(
        "-c", "--ncols", type=int, default=10, help="number of columns to select"
    )
    parser.add_argument(
        "-p",
        "--page-size",
        type=int,
        default=100,
        help="number of columns to to fetch per page (in psycopg > 3.2.10)",
    )
    parser.add_argument(
        "--repeat",
        type=int,
        default=10,
        help="number of times to repeat the test [default: %(default)s]",
    )
    parser.add_argument(
        "--no-schema", action="store_true", help="skip the schema creation"
    )
    parser.add_argument(
        "--relax", metavar="SEC", type=float, help="sleep for SEC between attempts"
    )

    opt = parser.parse_args()

    return opt


def setup_testtable():
    global opt

    print(f"creating table, cols: {opt.ncols}, rows: {opt.nrows}")

    with psycopg.connect(DSN) as conn:
        conn.execute("drop table if exists testtable")
        # fields = ", ".join(f"c{i} int" for i in range(opt.ncols))
        fields = ", ".join(f"c{i} text" for i in range(opt.ncols))
        conn.execute(
            f"create unlogged table testtable (id serial primary key, {fields})"
        )
        fields = ", ".join(f"c{i}" for i in range(opt.ncols))
        # values = ", ".join("10000" for i in range(opt.ncols))
        values = ", ".join("'aaaaaaaaaa'" for i in range(opt.ncols))
        conn.execute(
            f"insert into testtable ({fields}) select {values} from generate_series(1, %s)",
            [opt.nrows],
        )


if __name__ == "__main__":
    main()