"""Diagnostic: alternate installed JSON codecs per case in persistent workers."""
from __future__ import annotations
import argparse
from dataclasses import asdict
from hashlib import sha256
import json
import os
from pathlib import Path
from statistics import median
import subprocess
import sys
import time


def worker(source: Path, include: str) -> None:
    sys.path.insert(0, str(source))
    from bench.profiles import jsonx_throughput as profile
    from unstd.serde import jsonx
    from bench.core.timing import measure
    provenance = profile._provenance(source, jsonx, require_native=True)
    from functools import partial
    from types import SimpleNamespace
    import orjson
    cases = {}
    for label,value in [('float',1.5),('null',None),('int',7)]:
        cases['dumpb/'+label] = SimpleNamespace(name='dumpb/'+label,call=partial(jsonx.dumpb,value),reference=lambda v=value:json.dumps(v,ensure_ascii=False,separators=(',',':')).encode())
        cases['dumps/'+label] = SimpleNamespace(name='dumps/'+label,call=partial(jsonx.dumps,value),reference=partial(json.dumps,value))
        cases['raw-orjson/'+label] = SimpleNamespace(name='raw-orjson/'+label,call=partial(orjson.dumps,value),reference=lambda v=value:json.dumps(v,ensure_ascii=False,separators=(',',':')).encode())
    print(json.dumps({'provenance': provenance, 'cases': list(cases)}), flush=True)
    for line in sys.stdin:
        request = json.loads(line)
        case = cases[request['case']]
        expected = case.reference()
        if not profile._same(case.call(), expected):
            raise RuntimeError('incorrect output before timing: ' + case.name)
        start = time.time_ns()
        candidate = measure(case.call, rounds=request['rounds'])
        control = measure(case.reference, rounds=request['rounds'])
        if not profile._same(case.call(), expected):
            raise RuntimeError('incorrect output after timing: ' + case.name)
        print(json.dumps({'case': case.name, 'identical_output': True,
                          'started_unix_ns': start, 'finished_unix_ns': time.time_ns(),
                          'candidate': asdict(candidate), 'stdlib': asdict(control)}), flush=True)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--source', type=Path, required=True)
    parser.add_argument('--compare-source', type=Path)
    parser.add_argument('--harness', type=Path)
    parser.add_argument('--include', default='records-1k')
    parser.add_argument('--rounds', type=int, default=9)
    parser.add_argument('--repetitions', type=int, default=5)
    parser.add_argument('--max-slowdown', type=float, default=1.15)
    parser.add_argument('--output', type=Path)
    parser.add_argument('--worker', action='store_true')
    args = parser.parse_args()
    if args.worker:
        worker(args.source.resolve(), args.include)
        return
    assert args.compare_source and args.harness and args.output
    sources = {'before': args.compare_source.resolve(), 'after': args.source.resolve()}
    workers = {}
    provenance = {}
    samples = []
    try:
        for label, source in sources.items():
            env = dict(os.environ, PYTHONPATH=os.pathsep.join((str(source), str(args.harness.resolve()))), PYTHONHASHSEED='0')
            process = subprocess.Popen([sys.executable, '-u', str(Path(__file__).resolve()), '--worker', '--source', str(source), '--include', args.include], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, env=env)
            workers[label] = process
            ready = process.stdout.readline()
            if not ready:
                raise RuntimeError(process.stderr.read())
            provenance[label] = json.loads(ready)
        names = provenance['before']['cases']
        assert names == provenance['after']['cases'] and names
        for repetition in range(args.repetitions):
            ordered = names[::1 if repetition % 2 == 0 else -1]
            for index, name in enumerate(ordered):
                labels = ['before', 'after'][::1 if (repetition + index) % 2 == 0 else -1]
                for label in labels:
                    process = workers[label]
                    process.stdin.write(json.dumps({'case': name, 'rounds': args.rounds}) + '\n')
                    process.stdin.flush()
                    response = process.stdout.readline()
                    if not response:
                        raise RuntimeError(process.stderr.read())
                    samples.append(dict(label=label, repetition=repetition, **json.loads(response)))
            print(f'Measured per-case alternating repetition {repetition + 1}', file=sys.stderr)
        comparisons = {}
        for name in names:
            variants = {label: [row for row in samples if row['case'] == name and row['label'] == label] for label in sources}
            before = median(row['candidate']['per_op_ns'] for row in variants['before'])
            after = median(row['candidate']['per_op_ns'] for row in variants['after'])
            comparisons[name] = {'before_ns': before, 'after_ns': after, 'before_over_after': before / after,
                                 'before_stdlib_ns': median(row['stdlib']['per_op_ns'] for row in variants['before']),
                                 'after_stdlib_ns': median(row['stdlib']['per_op_ns'] for row in variants['after'])}
        report = {'driver_sha256': sha256(Path(__file__).read_bytes()).hexdigest(),
                  'protocol': 'Two persistent installed-root workers alternate per case; case order reverses each repetition; variant order alternates per case/repetition. Correctness and fixture creation are outside timing. No control normalization.',
                  'estimator': 'median across repetitions of minimum raw wall-clock batch means',
                  'rounds': args.rounds, 'repetitions': args.repetitions,
                  'provenance': provenance, 'comparisons': comparisons, 'samples': samples,
                  'max_slowdown': args.max_slowdown,
                  'regressions': [name for name, row in comparisons.items() if row['after_ns'] / row['before_ns'] > args.max_slowdown]}
        report['passed'] = not report['regressions']
        args.output.write_text(json.dumps(report, indent=2) + '\n')
    finally:
        for process in workers.values():
            if process.stdin:
                process.stdin.close()
            try:
                process.wait(timeout=10)
            except subprocess.TimeoutExpired:
                process.terminate()
                process.wait(timeout=10)

if __name__ == '__main__':
    main()
