"""Durable monitor and automatic packaging/upload after successful training."""
import json
import os
import re
import shutil
import signal
import subprocess
import sys
import time
from pathlib import Path

ROOT = Path('/data/realstyle2')
pid = int(sys.argv[1])
log = Path(sys.argv[2])
started = time.time()
target_epochs = 40
eta_reduced = False
while True:
    data = log.read_text(errors='replace')
    matches = re.findall(r'steps:\s+[^\r\n]*?\b(\d+)/(\d+)\s+\[([^\]]+)\](?:[^\r\n]*)', data)
    epochs = re.findall(r'epoch (\d+)/(\d+)', data)
    progress = matches[-1] if matches else None
    projected_seconds = None
    if progress and int(progress[0]) >= 100:
        rate = re.search(r'([\d.]+)s/it', progress[2])
        if rate:
            projected_seconds = float(rate[1]) * int(progress[1]) * target_epochs / 40
            if projected_seconds > 4 * 3600 and target_epochs > 30:
                target_epochs = 35 if float(rate[1]) * int(progress[1]) * 35 / 40 <= 4 * 3600 else 30
                eta_reduced = True
                (ROOT / 'metadata/eta_reduction.json').write_text(json.dumps({
                    'target_epochs': target_epochs, 'estimated_full_40_seconds': float(rate[1]) * int(progress[1]),
                    'at_step': int(progress[0]), 'reason': 'User requested a reduction when projected training exceeds 4 hours.'}, indent=2))
    stopped_at_limit = False
    if eta_reduced and epochs and int(epochs[-1][0]) > target_epochs:
        # Entering the next epoch proves the checkpoint and optimizer state were fully saved.
        checkpoint = ROOT / 'output' / f'realstyle2_anima-{target_epochs:06d}.safetensors'
        assert checkpoint.exists()
        os.kill(pid, signal.SIGTERM)
        for _ in range(30):
            try:
                if Path(f'/proc/{pid}/stat').read_text().split()[2] == 'Z':
                    break
            except FileNotFoundError:
                break
            time.sleep(1)
        shutil.copyfile(checkpoint, ROOT / 'output/realstyle2_anima.safetensors')
        import toml
        config_path = ROOT / 'configs/train.toml'
        config = toml.load(config_path)
        config['max_train_epochs'] = target_epochs
        config_path.write_text(toml.dumps(config))
        stopped_at_limit = True
    gpu = subprocess.check_output(['nvidia-smi', '--query-gpu=memory.used,utilization.gpu', '--format=csv,noheader'], text=True).strip()
    status = {'pid': pid, 'epoch': epochs[-1] if epochs else None, 'progress': progress,
              'gpu': gpu, 'updated_at': time.time(), 'elapsed_monitor_seconds': time.time()-started,
              'target_epochs': target_epochs, 'projected_training_seconds': projected_seconds,
              'saved_checkpoints': len(list((ROOT / 'output').glob('*.safetensors')))}
    (ROOT / 'status.json').write_text(json.dumps(status, indent=2))
    print(json.dumps(status), flush=True)
    try:
        os.kill(pid, 0)
        process_state = Path(f'/proc/{pid}/stat').read_text().split()[2]
        alive = process_state != 'Z'
    except (ProcessLookupError, FileNotFoundError):
        alive = False
    if stopped_at_limit:
        break
    if not alive:
        if 'model saved.' not in data or not (ROOT / 'output/realstyle2_anima.safetensors').is_file():
            raise RuntimeError('Training process exited before successful final save; inspect training log')
        break
    time.sleep(30)
print('Training complete. Validating, packaging, uploading.', flush=True)
with (ROOT / 'logs/upload.log').open('w') as out:
    subprocess.run(['/data/venv-training/bin/python', str(ROOT / 'scripts/package_upload.py')], stdout=out, stderr=subprocess.STDOUT, check=True)
print('UPLOAD COMPLETE', flush=True)