"""Sample live MySQL connection counts at high frequency.

WHY: pool pressure from a long-poll fleet is BURSTY. The route releases its
connection before sleeping and re-acquires on wake, so the interesting number is
the peak concurrent checkout during a wake burst, not the steady-state count. A
slow sampler misses it entirely.

Runs on the HOST (needs the docker CLI).

Usage:
    python sample_dbconns.py --seconds 120 --interval 0.25
"""
import argparse
import json
import statistics
import subprocess
import time

QUERY = ("SELECT COUNT(*), SUM(command<>'Sleep') "
         "FROM information_schema.processlist WHERE user='shopdb';")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--container', default='shopdb-flask-wilm-db-1')
    parser.add_argument('--seconds', type=int, default=120)
    parser.add_argument('--interval', type=float, default=0.25)
    arguments = parser.parse_args()

    total_samples = []
    active_samples = []
    deadline = time.monotonic() + arguments.seconds
    command = ('mysql -u root -p$MYSQL_ROOT_PASSWORD -N -B -e "%s"' % QUERY)

    while time.monotonic() < deadline:
        try:
            output = subprocess.check_output(
                ['docker', 'exec', arguments.container, 'sh', '-c', command],
                stderr=subprocess.DEVNULL, text=True).strip()
        except subprocess.CalledProcessError:
            break
        if not output:
            break
        parts = output.split()
        total_samples.append(int(parts[0]))
        active_samples.append(int(parts[1]) if parts[1] != 'NULL' else 0)
        time.sleep(arguments.interval)

    def summarize(values):
        if not values:
            return {'samples': 0}
        ordered = sorted(values)
        index = min(int(len(ordered) * 0.95), len(ordered) - 1)
        return {
            'samples': len(values),
            'median': statistics.median(values),
            'p95': ordered[index],
            'max': max(values),
        }

    print(json.dumps({
        'connections': summarize(total_samples),
        'active_queries': summarize(active_samples),
    }, indent=2))


if __name__ == '__main__':
    main()