# Timing a case: the interleaved rounds, the IO and instruction counts, the load measurement and the summary.
# Every answer is compared with the native one.

import os
import re
import signal
import statistics
import subprocess
import threading
import time
from collections import Counter

import benchlib
import iolib
from duckdb import Session

from . import runcfg
from .runcfg import FORCED, HERE, IO_HOW, MIN_RUNS, SCAN, die, log
from .server import Fresh, IDS, build, equal, equal_ties, plan_note, plan_shape, plan_text, psql, server_session, tie_keys


def expand(sql, case):
    """`{ids:TABLE:N}` is the list of the ids of the rows of TABLE with n = N, in the statement: a literal IN list of 50,000 or 500,000 values is not a thing to write in a file; `{param:TABLE:N}`
    is the same list as one int[] (a psql variable, :'ids_TABLE_N'::int[], sent once per session): the statement carries no 3 MB of SQL to parse and the planner still sees the array"""
    def ids(m):
        key = (case.schema(), m.group(1), m.group(2))
        if key not in IDS:
            IDS[key] = psql("SELECT string_agg(id::text, ',' ORDER BY id) FROM %s.%s WHERE n = %s" % key)
        return IDS[key]

    def param(m):
        name = 'ids_%s_%s' % (m.group(1), m.group(2))
        Session.VARS[name] = '{' + ids(m) + '}'
        return ":'%s'::int[]" % name
    sql = re.sub(r'\{param:(\w+):(\d+)\}', param, sql)
    return re.sub(r'\{ids:(\w+):(\d+)\}', ids, sql).replace('{rows}', str(runcfg.ROWS))


def tie_rows(case, native):
    """the rows of the ordering with the key values of a row (`--@ ties`, `{3}` the value of the third column of the row), from the lane of PostgreSQL's plan"""
    def rows(values):
        sql = re.sub(r'\{(\d+)\}', lambda m: values[int(m.group(1)) - 1].replace("'", "''"), case.sec['ties'])
        return native.run(sql)[1]
    return rows


def measure_query(case, lanes, rounds, runs, keep_plan=True):
    """rows of results: per lane the timed runs of the interleaved rounds, the answer of the warm-up compared with the native one; a plan other than the case's is an error
    unless keep_plan is false (the scale points: the planner's choice at each size is what the curve shows)"""
    ordered = case.opt.get('ordered', 'no')
    ties = tie_keys(ordered)
    sql = {lane: expand(case.statement(lane), case) for lane in lanes}
    notes, answers, times, errors, texts = {}, {}, {lane: [] for lane in lanes}, {}, {}
    for lane, s in lanes.items():
        text = texts[lane] = plan_text(s, sql[lane], lane != 'pg_duckdb')
        notes[lane] = plan_note(text)
        if lane == 'pg_duckdb' and 'DuckDBScan' not in text:
            errors[lane] = 'not supported: it runs on the PostgreSQL executor'
        ms, answers[lane], err, _ = s.run(sql[lane])   # the warm-up, whose answer is checked
        if err:
            errors.setdefault(lane, err[:110])
        if lane == 'sift':
            # the plan of the runs that are timed is the one after the first: where the first statement of a backend is planned another way (the cases cluster loads the
            # library at its start, so it is rare), the note says so
            again = plan_text(s, sql[lane])
            if plan_shape(again) != plan_shape(text):
                notes[lane] = '%s (the first plan of a backend was %s)' % (plan_note(again), notes[lane])
                text = texts[lane] = again
            if keep_plan and not re.search(case.plan, text, re.S):
                errors[lane] = 'PLAN LOST: expected /%s/, the plan is %s' % (case.plan, notes[lane])
            elif keep_plan and case.fallback and SCAN.search(text):
                errors[lane] = 'PLAN CHANGED: the planner chooses the scan now (%s): drop the fallback: of the plan column' % notes[lane]
            if case.fallback:
                notes['why'] = why(s, sql[lane])
    if FORCED[0] in texts and not SCAN.search(texts[FORCED[0]]) and FORCED[0] not in errors:
        notes[FORCED[0]] = 'no scan: ' + why(lanes[FORCED[0]], sql[FORCED[0]])   # the shape the scan cannot run: a feature gap, not a cost
    if 'native' in texts and 'sift' in texts and notes['sift'] == notes['native'] and plan_shape(texts['sift']) != plan_shape(texts['native']):
        notes['sift'] += ' (a different plan below it)'   # the first scan is the same: a block must not call two different plans "the same plan"
    ref = answers.get('native')
    for lane in lanes:
        if lane not in errors and lane != 'native' and ref is not None and not (equal_ties(ref, answers[lane], ties, tie_rows(case, lanes['native'])) if ties else equal(ref, answers[lane], ordered == 'yes')):
            errors[lane] = 'RESULT DIFFERS from native (%d rows against %d)' % (len(answers[lane]), len(ref))
    for r in range(rounds):
        order = [l for l in lanes if l not in errors]
        order = order[r % len(order):] + order[:r % len(order)] if order else order
        for lane in order:
            for _ in range(runs):
                if sum(times[lane]) / 1000 > runcfg.LANE_BUDGET and len(times[lane]) >= MIN_RUNS:   # a lane of seconds a run: the median of a few runs, not of fifteen
                    break
                times[lane].append(lanes[lane].run(sql[lane])[0])
    return times, notes, errors, texts


INDEX_IN_PLAN = re.compile(r'(?:Index(?: Only)? Scan(?: Backward)? using|Bitmap Index Scan on) (\w+)|Index: (\w+)')


def io_load(session, case, lane, sql, text, mode):
    """what a lane reads, whatever the storage (steps/iolib.py): the IO of the first run after the case's tables were emptied from shared buffers and the OS cache (cold), then of the run
    after it (warm), and the bytes of the indexes the plan reads (a btree, the column index of a scan). MODE warm (and off) skips the eviction: one warm run. Without pg_buffercache the one run is warm.
    Returns ([measurements], index names, index bytes)"""
    s = session.persistent() if isinstance(session, Fresh) else session
    try:
        schema = case.schema() + '_base' if lane == 'base' else case.schema()
        for name, settings in case.variants:   # a lane that reads the copy in another schema (`btree:search_path = c_ID__btree`) empties that copy, not the case's own tables
            m = re.search(r'search_path = (\w+)', settings)
            if name == lane and m:
                schema = m.group(1)
        tables = iolib.schema_tables(s, schema)
        ios = [iolib.measure(s, sql, tables if mode == 'cold' else None)]
        if ios[0]['cache'] == 'cold':
            IO_HOW.add(ios[0]['how'])
            ios.append(iolib.measure(s, sql))
        elif mode == 'cold':
            log('%s on %s: warm only (%s)' % (case.id, lane, ios[0]['how']))
        names = sorted({a or b for a, b in INDEX_IN_PLAN.findall(text)})
        size = sum(int(s.run("SELECT pg_relation_size('%s')" % n)[1][0]) for n in names) if names else 0
        return ios, names, size
    except RuntimeError as e:   # a statement the lane cannot EXPLAIN or finish in the timeout has no IO rows; the run goes on
        log('%s on %s: no IO rows (%s)' % (case.id, lane, str(e)[:100]))
        return [], [], 0
    finally:
        if s is not session:
            s.close()


def why(session, sql):
    """sift.explain_why of the statement (planned, not run): the first reason each arm of the aggregate pushdown gave up"""
    out = session.run("SELECT arm || ' ' || slug || ': ' || reason FROM sift.explain_why($q$%s$q$) WHERE NOT used" % sql.rstrip(';'))[1]
    return '; '.join(out)[:300] or 'sift.explain_why names no reason'


PERF = os.path.join(HERE, 'perf.sh')


class PerfFailed(Exception):
    """perf did not give the instructions of a statement in any of its attempts; the message is the reason of each"""


def perf_paranoid():
    """the perf_event_paranoid setting for a message about a perf that failed; where the host has none (not Linux) the message says so"""
    try:
        with open('/proc/sys/kernel/perf_event_paranoid') as f:
            return f.read().strip()
    except OSError:
        return 'not readable here'


def count_once(session, sql, pid):
    """(millions of user instructions of one run of SQL by the backend PID, None) or (None, why): steps/perf.sh, perf stat on the backend with no parallel workers"""
    p = subprocess.Popen(['bash', PERF, pid], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, start_new_session=True)
    time.sleep(1)   # perf.sh traps SIGTERM only once bash has started: on a loaded host 0.3 s was too short and the signal killed bash, leaving perf to count to the end of the backend
    session.run(sql)
    p.send_signal(signal.SIGTERM)
    try:
        out, err = p.communicate(timeout=60)   # a SIGTERM that lands before perf.sh traps it leaves perf counting for ever
    except subprocess.TimeoutExpired:
        os.killpg(p.pid, signal.SIGKILL)
        return None, 'perf did not end 60 s after the SIGTERM (it reached perf.sh before its trap)'
    fields = out.split()
    if len(fields) == 3 and fields[1].isdigit():
        return int(fields[1]) / 1e6, None
    return None, 'perf counted "%s" (perf_event_paranoid %s, exit %s): %s' % (' '.join(fields), perf_paranoid(), p.returncode, err.strip()[-200:] or 'no message')


def minstr(session, sql, runs=3, tries=3):
    """user instructions (millions) of the backend of the session for one run of the statement, median of `runs`; a count that perf could not give is tried again up to TRIES times, then
    PerfFailed says why each attempt failed (the count does not depend on the load of the host, so a retry is the same measurement)"""
    session.run('SET max_parallel_workers_per_gather = 0')
    pid = session.run('SELECT pg_backend_pid()')[1][0]
    session.run(sql)
    counts = []
    for _ in range(runs):
        whys = []
        for _ in range(tries):
            n, why = count_once(session, sql, pid)
            if why is None:
                counts.append(n)
                break
            whys.append(why)
        else:
            raise PerfFailed('; '.join('attempt %d: %s' % (k + 1, w) for k, w in enumerate(whys)))
    return statistics.median(counts)


def stats(ms):
    return dict(value='%.3f' % statistics.median(ms), min='%.3f' % min(ms), max='%.3f' % max(ms), n=len(ms))


def wal_lsn(s):
    return int(s.run("SELECT pg_wal_lsn_diff(pg_current_wal_lsn(), '0/0')::bigint")[1][0])


def cpu_ns(pid):
    with open('/proc/%s/schedstat' % pid) as f:
        return int(f.read().split()[0])


def wait_sampler(pid):
    """the kernel wait (state:function) of the backend every 2 ms while it does not run, until stop() returns: where the time of a load beyond its CPU went"""
    seen, done = Counter(), threading.Event()

    def sample():
        while not done.wait(0.002):
            try:
                with open('/proc/%s/stat' % pid) as f:
                    state = f.read().rsplit(')', 1)[1].split()[0]
                if state in ('S', 'D'):
                    with open('/proc/%s/wchan' % pid) as f:
                        seen[state + ':' + f.read()] += 1
            except OSError:
                return
    thread = threading.Thread(target=sample)
    thread.start()

    def stop():
        done.set()
        thread.join()
        return seen
    return stop


def record_start(pid, path, perf='perf'):
    """perf record of the user cycles of the backend PID into PATH (the profile of a load phase: `--profile DIR`); record_stop ends it"""
    return subprocess.Popen([perf, 'record', '-q', '-F', '2000', '-e', 'cycles:u', '-p', pid, '-o', path], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True, start_new_session=True)


def record_stop(proc, path, perf='perf', top=30):
    """end the record and write the PATH's symbols by self cycles to PATH.txt (the first TOP); die with perf's message when it recorded nothing"""
    proc.send_signal(signal.SIGTERM)   # perf record finishes its file on SIGTERM too; a job started in the background has SIGINT ignored
    try:
        err = proc.communicate(timeout=60)[1]
    except subprocess.TimeoutExpired:
        os.killpg(proc.pid, signal.SIGKILL)
        die('perf record did not end 60 s after the SIGTERM (%s)' % path)
    if proc.returncode not in (0, -signal.SIGTERM) or not os.path.exists(path):   # perf finishes its file, then raises the signal again
        die('perf record failed (exit %s, perf_event_paranoid %s): %s' % (proc.returncode, perf_paranoid(), err.strip()[-300:] or 'no message'))
    report = subprocess.run([perf, 'report', '-i', path, '--no-children', '--stdio'], capture_output=True, text=True).stdout
    with open(path + '.txt', 'w') as f:
        f.write('\n'.join([l for l in report.split('\n') if l.strip() and not l.startswith('#')][:top]) + '\n')


def load_settings(case_id, workers, sets, serial):
    """the statements that open the session of a pass of a load case: the schema of the case, the parallel maintenance workers of the run, the settings of the variant and, in the serial pass of
    the instructions, workers 0 as the last one so it overrides the variant's own (perf follows the children of the backend, and the workers of a VACUUM, a build or a merge are children of the postmaster)"""
    return ['SET search_path = c_%s, public' % case_id, 'SET max_parallel_maintenance_workers = %s' % workers] + sets + (['SET max_parallel_maintenance_workers = 0'] if serial else [])


def load_passes(instr):
    """the passes of one lane in one round: the timed one, and with INSTR a serial one for the instructions"""
    return (False, True) if instr else (False,)


def measure_load(case, csv, rounds, rows, instr=False, profile=None):
    """a load case: per lane and round one fresh table (`--@ prepare`, untimed, with {am} and {csv}), the statements of `--@ load` timed with its WAL and the backend's CPU; metrics wal_bytes (per row), ms, cpu_ms, index_mb (the indexes other than the primary key, at the end), and with INSTR the user instructions of the whole load in a second, serial run of it (minstr: the backend alone, `max_parallel_maintenance_workers = 0` on every lane, the settings of workers of a variant overridden: the workers of a VACUUM, a build or a merge are children of the postmaster and perf does not follow them)"""
    lanes = {'sift': ('sift_table', []), 'native': ('heap', [])}
    for name, settings in case.variants:
        lanes[name] = ('sift_table', ['SET %s' % s for s in settings.split(';')])
    out, errors, plans = {lane: {'wal_bytes': [], 'ms': [], 'cpu_ms': [], 'io_extend_mb': [], 'io_write_mb': [], 'index_mb': []} | ({'minstr': []} if instr else {}) for lane in lanes}, {}, {}
    for r in range(rounds):
        names = list(lanes)
        for lane in names[r % len(names):] + names[:r % len(names)]:
            am, sets = lanes[lane]
            for serial in load_passes(instr):   # the instructions in a pass of their own with no parallel workers: perf follows the children of the backend, and the workers of a VACUUM, a build or a merge are children of the postmaster
                s = server_session(600)
                for st in load_settings(case.id, os.environ.get('CASES_WORKERS', '2'), sets, serial):
                    s.run(st)
                for stmt in re.split(r';[ \t]*\n', case.sec['prepare'].replace('{am}', am).replace('{csv}', csv).replace('{rows}', str(rows))):   # the statements of the preparation, untimed (a VACUUM among them cannot share a transaction)
                    err = s.run(stmt)[2]
                    if err:
                        die('%s: the preparation failed on %s: %s' % (case.id, lane, err))
                pid = s.run('SELECT pg_backend_pid()')[1][0]
                s.run('CHECKPOINT')
                io0 = iolib.snapshot(s)
                lsn, cpu = wal_lsn(s), cpu_ns(pid)
                counter = subprocess.Popen(['bash', PERF, pid], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, start_new_session=True) if serial else None
                recorded = os.path.join(profile, '%s.%s.data' % (case.id, lane)) if profile and not serial else None
                recorder = record_start(pid, recorded) if profile and not serial else None
                if counter or recorder:
                    time.sleep(1)
                ms = 0
                stop = wait_sampler(pid)
                for stmt in re.split(r';[ \t]*\n', case.sec.get('load_sift', case.sec['load']) if am == 'sift_table' else case.sec['load']):   # the statements of the load, their times added
                    took, _, err, _ = s.run(stmt.replace('{csv}', csv).replace('{rows}', str(rows)))
                    if err:
                        die('%s: the load failed on %s: %s' % (case.id, lane, err))
                    ms += took
                waits = stop()
                if recorder:
                    record_stop(recorder, recorded)
                cpu, lsn = cpu_ns(pid) - cpu, wal_lsn(s) - lsn
                if ms - cpu / 1e6 > max(20, ms / 5):
                    log('%s %s: %.0f of %.0f ms off the CPU, 2 ms samples of the kernel wait: %s' % (case.id, lane, ms - cpu / 1e6, ms, ', '.join('%s %d' % kv for kv in waits.most_common(4))))
                if counter:
                    counter.send_signal(signal.SIGTERM)
                    fields = counter.communicate()[0].split()
                    if len(fields) == 3 and fields[1].isdigit():
                        out[lane]['minstr'].append(int(fields[1]) / 1e6)
                if serial:
                    s.close()
                    build(case, rows)
                    continue
                for metric, mb in iolib.load_metrics(iolib.delta(io0, iolib.snapshot(s))).items():
                    out[lane][metric].append(mb)
                checked = '\n'.join(s.run(case.sec['check'])[1]) if lane == 'sift' else ''
                plans[lane] = checked.replace('\n', ', ')
                if lane == 'sift' and not re.search(case.plan, checked, re.M):
                    errors[lane] = 'PLAN LOST: the load did not write what /%s/ describes' % case.plan
                size = s.run("SELECT sum(pg_relation_size(i.indexrelid)) / 1048576.0 FROM pg_index i JOIN pg_class c ON c.oid = i.indrelid JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = 'c_%s' AND NOT i.indisprimary" % case.id)[1][0]
                if size:   # a case with no index besides its primary key has no figure
                    out[lane]['index_mb'].append(float(size))
                out[lane]['wal_bytes'].append(lsn / rows)
                out[lane]['ms'].append(ms)
                out[lane]['cpu_ms'].append(cpu / 1e6)
                s.close()
                build(case, rows)   # an empty schema again
    return out, plans, errors


def host():
    return benchlib.host_label()


def summarize(results, cases):
    """the table printed after a run and the loss lines: where sift is slower than native, plainly"""
    lines, losses = ['%-26s %-10s %10s %10s %8s  %s' % ('case', 'metric', 'native', 'sift', 'speedup', 'plan')], []
    for c, m in [(c, m) for c in cases for m in c.metrics]:
        cell = lambda arm: next((r for r in results if r['case'] == c.id and r['arm'] == arm and r['metric'] == m), None)  # noqa: E731
        n, s = cell('native'), cell('sift')
        if not (n and s and 'NA' not in (n['value'], s['value']) and 'error' not in (n['value'], s['value'])):
            lines.append('%-26s %-10s %10s %10s %8s  %s' % (c.id, m, n['value'] if n else '-', s['value'] if s else '-', '-', (s or {}).get('note', '')[:60]))
            continue
        if c.fallback:   # PostgreSQL's plan against itself has no ratio: the lane with the scan forced is the comparison
            f = cell(FORCED[0])
            ok = f and f['value'] not in ('NA', 'error', '')
            lines.append('%-26s %-10s %10s %10s %8s  %s' % (c.id, m, n['value'], f['value'] if ok else '-', '%.1fx' % (float(n['value']) / float(f['value'])) if ok and float(f['value']) else '-', 'scan forced: ' + f['note'][:46] if f else ''))
            continue
        speed = float(n['value']) / float(s['value']) if float(s['value']) else float('inf')
        lines.append('%-26s %-10s %10s %10s %7.1fx  %s' % (c.id, m, n['value'], s['value'], speed, s['note'][:60]))
        if speed < 0.95:
            losses.append('%s: sift %s %s against native %s (%.1fx slower)' % (c.id, s['value'], m, n['value'], 1 / speed))
    if losses:
        lines += ['', 'slower than native on %d of %d numbers:' % (len(losses), sum(len(c.metrics) for c in cases))] + ['  ' + l for l in losses]
    return '\n'.join(lines)