#!/usr/bin/env python3
import subprocess, json, re
from datetime import datetime, timezone

DATA_FILE = '/opt/guanghulab-repo/homepage/training-status.json'
SSH = ['sshpass','-p','HkM43lFVUIsc','ssh','-o','StrictHostKeyChecking=no',
       '-o','ConnectTimeout=10','root@connect.westd.seetacloud.com','-p','23647']

try:
    r = subprocess.run(SSH + ['tail','-2','/root/autodl-tmp/train_mother.log'],
                      capture_output=True, text=True, timeout=25).stdout
    sm = re.search(r'(\d+)/12243', r)
    lm = re.search(r"'loss':\s*'([\d.]+)'", r)
    step = int(sm.group(1)) if sm else 0
    loss = lm.group(1) if lm else '--'
except:
    step, loss = 0, '--'

data = {'step': step, 'total': 12243, 'loss': loss,
        'updated': datetime.now(timezone.utc).isoformat()+'Z'}

with open(DATA_FILE, 'w') as f:
    json.dump(data, f)
print(f"Step {step} / 12243, Loss: {loss}")