#!/usr/bin/env python3
import datetime
import hashlib
import json
import os
from pathlib import Path
import shutil
import signal
import subprocess
import tempfile
import time

ROOT = Path('/home/reavus/projects/spt-core/spt-core')
TREE = ROOT / '.worktrees/304-release-note-s4'
TARGET = TREE / 'target'
PROOF = Path(__file__).resolve().parent
SHA = '6c89e8f7545db54772ea5686b4d59718573d1f08'
ALLOWED = {'CHANGELOG.md', 'docs-site/src/changelog.md'}
FLOOR = 32 * 1024**3
HELPER = ROOT / '.spt/preserved/308-registry-process-lock/linux-validation/xtask-b8482445'
ENV = {k: v for k, v in os.environ.items() if not k.startswith(('OWL_', 'SPT_'))}
for k in ('CARGO_ENCODED_RUSTFLAGS', 'NEXTEST_TEST_THREADS', 'NEXTEST_RETRIES'):
    ENV.pop(k, None)
ENV.update(PATH='/home/reavus/.cargo/bin:' + ENV['PATH'], CARGO_TARGET_DIR=str(TARGET),
           CARGO_BUILD_JOBS='2', CARGO_INCREMENTAL='0', RUSTFLAGS='-C link-arg=-fuse-ld=mold',
           CARGO_TERM_COLOR='never', CI='true')
TEMP = Path(tempfile.mkdtemp(prefix='spt-304-s4-docs-', dir='/tmp'))
ENV.update(TEMP=str(TEMP), TMP=str(TEMP), TMPDIR=str(TEMP), SPT_HOME=str(TEMP / 'home'))

def utc():
    return datetime.datetime.now(datetime.timezone.utc).isoformat()

def save(name, value):
    (PROOF / name).write_text(json.dumps(value, indent=2) + '\n')

def census():
    active, own = [], []
    for p in Path('/proc').iterdir():
        if not p.name.isdigit():
            continue
        try:
            name = (p / 'comm').read_text().strip()
            exe = os.readlink(p / 'exe')
            row = {'pid': int(p.name), 'name': name, 'exe': exe,
                   'command': (p / 'cmdline').read_bytes().replace(b'\0', b' ').decode(errors='replace'),
                   'stat': (p / 'stat').read_text()}
            if any(x in name for x in ('cargo', 'rustc', 'nextest', 'Runner.Worker')) or '/deps/' in exe:
                active.append(row)
            if exe.startswith(str(TARGET) + '/'):
                own.append(row)
        except OSError:
            pass
    return {'utc': utc(), 'free_bytes': shutil.disk_usage(TREE).free,
            'active': active, 'own_survivors': own}

def source():
    head = subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=TREE, text=True).strip()
    changed = subprocess.check_output(['git', 'diff', '--name-only'], cwd=TREE, text=True).splitlines()
    new = subprocess.check_output(['git', 'ls-files', '--others', '--exclude-standard'], cwd=TREE, text=True).splitlines()
    assert head == SHA and set(changed) <= ALLOWED and not new, (head, changed, new)
    return {'base': head, 'changed': changed,
            'hashes': {name: hashlib.sha256((TREE / name).read_bytes()).hexdigest() for name in ALLOWED}}

def admit(label):
    state = {**census(), 'source': source()}
    save(label + '-admission.json', state)
    assert (label == 'pool-release' or state['free_bytes'] >= FLOOR) and not state['active'] and not state['own_survivors'], state

def run(label, argv):
    admit(label)
    start = utc()
    clock = time.monotonic()
    samples = []
    crossed = False
    log = PROOF / (label + '.log')
    print('BEGIN ' + label + ' ' + json.dumps(argv), flush=True)
    with log.open('x') as stream:
        child = subprocess.Popen(argv, cwd=TREE, env=ENV, stdout=stream,
                                 stderr=subprocess.STDOUT, start_new_session=True)
        while child.poll() is None:
            free = shutil.disk_usage(TREE).free
            samples.append({'utc': utc(), 'free_bytes': free})
            if free < FLOOR and label != 'pool-release':
                crossed = True
                os.killpg(child.pid, signal.SIGKILL)
                child.wait()
                break
            time.sleep(1)
    row = {'argv': argv, 'cwd': str(TREE), 'start_utc': start, 'end_utc': utc(),
           'elapsed_s': time.monotonic() - clock, 'exit': child.returncode,
           'floor_crossed': crossed, 'samples': samples,
           'log_sha256': hashlib.sha256(log.read_bytes()).hexdigest(), 'source_after': source()}
    save(label + '-receipt.json', row)
    print('END ' + label + ' ' + json.dumps({k: v for k, v in row.items() if k != 'samples'}), flush=True)
    assert not crossed and child.returncode == 0, row
    return row

assert not (PROOF / 'receipt.json').exists()
claimed, success = False, False
rows = []
try:
    save('environment.json', {'keys': sorted(ENV), 'owl_spt_values': {k: v for k, v in ENV.items() if k.startswith(('OWL_', 'SPT_'))},
                             'target': str(TARGET), 'jobs': ENV['CARGO_BUILD_JOBS'], 'rustflags': ENV['RUSTFLAGS'],
                             'rustfmt_invoked': False})
    admit('start')
    assert not any(TARGET.iterdir()), 'S4 target was not cold'
    claim = run('pool-claim', [str(HELPER), 'pool-claim', '--pool', str(TARGET), '--label', 'todlando-304-s4-docs'])
    claimed = True
    rows.append(run('xtask-gen', ['cargo', 'run', '-p', 'xtask', '--', 'gen']))
    assert set(source()['changed']) == ALLOWED
    rows.append(run('xtask-check', ['cargo', 'run', '-p', 'xtask', '--', 'check']))
    rows.append(run('trace-check', ['/home/reavus/.local/bin/traceable-reqs', 'check', '--json']))
    trace = json.loads((PROOF / 'trace-check.log').read_text())
    save('trace-summary.json', trace['summary'])
    success = True
except BaseException as error:
    save('driver-error.json', {'error': repr(error), 'census': census(), 'source': source()})
    raise
finally:
    release = None
    if claimed:
        release = run('pool-release', [str(HELPER), 'pool-release', '--pool', str(TARGET)])
    final = census()
    success = success and not final['active'] and not final['own_survivors']
    save('receipt.json', {'base': SHA, 'success': success, 'phases': rows, 'pool_release': release,
                          'final': final, 'source': source(), 'rustfmt_invoked': False})
    print('S4_DOCS_DONE ' + json.dumps({'success': success, 'final': final}), flush=True)
raise SystemExit(0 if success else 1)