"""One launch, fresh resource guards, only owned scratch/process-group cleanup."""
import hashlib
import json
import os
from pathlib import Path
import shutil
import signal
import subprocess
import tempfile
import time
from datetime import datetime, timezone

ROOT = Path('/home/frank14f/CelerisLab')
OUT = Path(__file__).resolve().parent
PYTHON = '/home/frank14f/anaconda3/envs/pycuda_3_10/bin/python'
UUID = 'GPU-e8853ae7-a302-6a17-bcfa-aa767e9ed20c'


def command(args):
    return subprocess.check_output(args, text=True, timeout=10)


def gpu():
    return command(['nvidia-smi', '--query-gpu=index,uuid,name,memory.used,memory.free,utilization.gpu', '--format=csv,noheader,nounits'])


def apps():
    return command(['nvidia-smi', '--query-compute-apps=gpu_uuid,pid,process_name,used_memory', '--format=csv,noheader,nounits'])


def ram():
    for line in Path('/proc/meminfo').read_text().splitlines():
        if line.startswith('MemAvailable:'):
            return int(line.split()[1])*1024
    raise RuntimeError('MemAvailable unavailable')


def disk(paths):
    return sum(p.stat().st_size for root in paths for p in root.rglob('*') if p.is_file())


def group_rss(pgid):
    total = 0
    for proc in Path('/proc').iterdir():
        if not proc.name.isdigit():
            continue
        try:
            fields = (proc/'stat').read_text().rsplit(') ', 1)[1].split()
            if int(fields[2]) == pgid:
                for line in (proc/'status').read_text().splitlines():
                    if line.startswith('VmRSS:'):
                        total += int(line.split()[1])*1024
        except (OSError, ValueError, IndexError):
            pass
    return total


def main():
    assert not (OUT/'gpu_run.json').exists(), 'Do not repeat an existing launch'
    data = dict(start_utc=datetime.now(timezone.utc).isoformat(), command=[PYTHON, str(OUT/'validate.py')],
                cwd=str(ROOT), physical_gpu_uuid=UUID, note='Sampled guards, not an OS hard isolation; no fixed RSS/AS cap.')
    scratch = None
    process = None
    started = time.monotonic()
    code = 1
    try:
        frozen = json.loads((OUT/'frozen.json').read_text())
        for path, digest in frozen['hashes'].items():
            assert hashlib.sha256((ROOT/path).read_bytes()).hexdigest() == digest, path
        data['preflight_gpu'], data['preflight_apps'], data['preflight_ram_bytes'] = gpu(), apps(), ram()
        fields = next(line.split(', ') for line in data['preflight_gpu'].splitlines() if UUID in line)
        used, free, util = map(int, fields[-3:])
        assert fields[0] == '0' and used <= 16 and free >= 2048 and util == 0
        assert UUID not in data['preflight_apps'] and data['preflight_ram_bytes'] >= 4*1024**3
        limits = dict(wall_seconds=180, temp_bytes=2_000_000_000,
                      ram_reserve_bytes=max(4*1024**3, data['preflight_ram_bytes']//10),
                      vram_reserve_MiB=max(512, free//10))
        data['limits'] = limits
        scratch = Path(tempfile.mkdtemp(prefix='celeris-joint24-'))
        env = dict(PYTHONPATH=str(ROOT/'src'), PYTHONNOUSERSITE='1', PYTHONDONTWRITEBYTECODE='1',
                   OPENBLAS_NUM_THREADS='1', CUDA_VISIBLE_DEVICES=UUID,
                   CUDA_CACHE_DISABLE='1', TMPDIR=str(scratch))
        data['env'] = env
        data['frozen_sha256'] = hashlib.sha256((OUT/'frozen.json').read_bytes()).hexdigest()
        data['peak_group_rss_bytes'] = 0
        data['peak_temp_bytes'] = disk([OUT,scratch])
        data['min_ram_bytes'] = data['preflight_ram_bytes']
        data['min_vram_free_MiB'] = free
        data['termination_reason'] = None
        with (OUT/'stdout.txt').open('w') as output:
            process = subprocess.Popen(data['command'], cwd=ROOT, env={**os.environ, **env},
                                       stdout=output, stderr=subprocess.STDOUT, start_new_session=True)
            while process.poll() is None:
                available, occupied = ram(), disk([OUT,scratch])
                free = int(next(line.split(', ')[-2] for line in gpu().splitlines() if UUID in line))
                data['peak_group_rss_bytes'] = max(data['peak_group_rss_bytes'], group_rss(process.pid))
                data['peak_temp_bytes'] = max(data['peak_temp_bytes'], occupied)
                data['min_ram_bytes'] = min(data['min_ram_bytes'], available)
                data['min_vram_free_MiB'] = min(data['min_vram_free_MiB'], free)
                if (time.monotonic()-started > limits['wall_seconds'] or occupied > limits['temp_bytes']
                        or available < limits['ram_reserve_bytes'] or free < limits['vram_reserve_MiB']):
                    data['termination_reason'] = 'Time/disk/RAM/VRAM guard triggered'
                    os.killpg(process.pid, signal.SIGTERM)
                    try:
                        process.wait(timeout=5)
                    except subprocess.TimeoutExpired:
                        os.killpg(process.pid, signal.SIGKILL)
                    break
                time.sleep(.5)
            code = process.wait(timeout=10)
        data['postflight_gpu'], data['postflight_apps'] = gpu(), apps()
    except BaseException as error:
        data['runner_error'] = repr(error)
        if process is not None and process.poll() is None:
            os.killpg(process.pid, signal.SIGKILL)
            process.wait(timeout=10)
        raise
    finally:
        data['exit_code'] = code
        data['elapsed_seconds'] = time.monotonic()-started
        if scratch is not None:
            shutil.rmtree(scratch)
        data['scratch_removed'] = scratch is None or not scratch.exists()
        (OUT/'gpu_run.json').write_text(json.dumps(data, indent=2)+'\n')
        print(json.dumps(data, indent=2), flush=True)
    raise SystemExit(code)


if __name__ == '__main__':
    main()