#!/usr/bin/env python3
"""Bounded watchdog for ONLY the separately authorized Round5V4 restore container."""
import argparse
import json
import os
from pathlib import Path
import shutil
import subprocess
import time

NAME = 'railwaysystem-qa-restore-round5v4'
IMAGE = 'sha256:540fa4d0c8f365c8d363a801404e93f8401c7bc5035ecd7b621f0dba862f2515'
ROOT = Path('/opt/railwaysystem-qa-restore-round5v4')
LIMIT = 4 * 1024 ** 3  # stop below the approved 5GiB budget; bind has NO hard quota
RESERVE = 100 * 1024 ** 3


def command(*args):
    return subprocess.check_output(args, text=True, timeout=20).strip()


def inspect(container_id):
    item = json.loads(command('docker', 'inspect', container_id))[0]
    if item['Id'] != container_id or item['Name'] != '/' + NAME or item['Image'] != IMAGE:
        raise RuntimeError('Container identity mismatch; no unrelated process will be signaled')
    if item['Config'].get('Labels', {}).get('railway.task') != 'round5v4':
        raise RuntimeError('Missing task ownership label')
    mounts = [m for m in item['Mounts'] if m['Destination'] == '/home/kingbase/userdata']
    if len(mounts) != 1 or Path(mounts[0]['Source']).resolve() != ROOT / 'data':
        raise RuntimeError('Unexpected data mount')
    return item


def used_bytes(path):
    return int(command('du', '-sk', str(path)).split()[0]) * 1024


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--once', action='store_true')
    args = parser.parse_args()
    if ROOT.is_symlink() or ROOT.resolve() != ROOT or (ROOT / 'data').is_symlink():
        raise RuntimeError('Task paths must not be symlinks')
    container_id = (ROOT / 'container.id').read_text().strip()
    ever_running = False
    while True:
        item = inspect(container_id)
        running = item['State']['Running']
        ever_running = ever_running or running
        upper = Path(item['GraphDriver']['Data']['UpperDir']).resolve()
        if not str(upper).startswith('/var/lib/docker/overlay2/') or upper.name != 'diff':
            raise RuntimeError('Unexpected writable-layer location')
        used = used_bytes(ROOT) + used_bytes(upper)
        log_path = item.get('LogPath')
        used += Path(log_path).stat().st_size if log_path and Path(log_path).exists() else 0
        free = shutil.disk_usage(ROOT).free
        print(json.dumps({'time': time.time(), 'running': running, 'used_bytes': used,
                          'stop_threshold_bytes': LIMIT, 'host_free_bytes': free}), flush=True)
        if used >= LIMIT or free < RESERVE:
            # Revalidate immutable task ownership immediately before stopping.
            inspect(container_id)
            subprocess.run(['docker', 'stop', '--time', '15', container_id], check=True, timeout=25,
                           stdout=subprocess.DEVNULL)
            print('STOPPED_TASK_CONTAINER_STORAGE_THRESHOLD', flush=True)
            return
        if args.once or (ever_running and not running):
            return
        time.sleep(2)


if __name__ == '__main__':
    main()