#!/usr/bin/env python3
"""
fetch_train.py v4 — 蒸馏+GPU实时进度同步 (FIXED: phase/epoch/step detection)
"""
import subprocess, json, re
from datetime import datetime, timezone

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

def ssh(cmd_list):
    try:
        r = subprocess.run(SSH_BASE + cmd_list, capture_output=True, text=True, timeout=30)
        return r.stdout.strip()
    except:
        return ''

def get_distill_status():
    # Use tail -300 to ensure we catch epoch print + loss lines
    raw = ssh(['tail','-300','/root/autodl-tmp/distill_mother.log'])
    if not raw:
        return {}, 'no_log'
    lines = raw.replace('\r', '\n').split('\n')

    # Phase detection: look for tqdm progress bar pattern EpX: X%|
    if 'ALL DONE' in raw or 'ALL VERIFIED' in raw:
        return {'phase': 'completed'}, 'completed'
    
    # Check if tqdm is running (EpX: XX%| pattern)
    has_progress = False
    for line in lines:
        if re.search(r'Ep\d+:\s+\d+%\|', line):
            has_progress = True
            break
    
    if has_progress:
        phase = 'training'
    elif 'Train' in raw or 'Training' in raw:
        phase = 'training'
    elif 'Load models' in raw or 'Loading weights' in raw:
        phase = 'loading_models'
    elif 'Tokenize' in raw:
        phase = 'tokenizing'
    elif 'Save final' in raw or 'Upload' in raw:
        phase = 'saving'
    else:
        phase = 'running'

    # Epoch: try "EpX" from tqdm first, then "Epoch X/Y" from print
    epoch = None; total_epoch = 3
    for line in reversed(lines):
        m = re.search(r'Ep(\d+):', line)
        if m:
            epoch = int(m.group(1))
            break
    if epoch is None:
        for line in reversed(lines):
            m = re.search(r'Epoch\s+(\d+)/(\d+)', line)
            if m:
                epoch, total_epoch = int(m.group(1)), int(m.group(2))
                break

    # Step & total: from tqdm progress bar  step/15781
    step = None; total_steps = None
    for line in reversed(lines):
        m = re.search(r'(\d+)/(\d+)\s+\[', line)
        if m:
            step, total_steps = int(m.group(1)), int(m.group(2))
            break

    # Loss: from "step=X loss=Y.ZZZZ" lines
    loss = '--'
    for line in reversed(lines):
        m = re.search(r"loss=([0-9.]+)", line)
        if m:
            loss = m.group(1)
            break

    # ETA: from tqdm progress <HH:MM:SS>
    eta = '--'
    for line in reversed(lines):
        m = re.search(r'<([0-9]+:[0-9]+:[0-9]+)', line)
        if m:
            eta = m.group(1)
            break

    return {'phase': phase, 'epoch': epoch, 'total_epoch': total_epoch,
            'step': step, 'total_steps': total_steps, 'loss': loss, 'eta': eta}, phase

def get_gpu_status():
    raw = ssh(['nvidia-smi','--query-gpu=temperature.gpu,memory.used,memory.total,utilization.gpu,utilization.memory',
               '--format=csv,noheader'])
    parts = raw.split(', ')
    if len(parts) >= 5:
        return {'temp_c': parts[0].strip(),
                'mem_used_gb': round(int(parts[1].strip().split()[0])/1024, 1),
                'mem_total_gb': round(int(parts[2].strip().split()[0])/1024, 1),
                'gpu_util_pct': parts[3].strip().split()[0],
                'mem_util_pct': parts[4].strip().split()[0]}
    return {}

def main():
    info = {'mode': 'distill_shuangyan'}
    status, phase = get_distill_status()
    info.update(status)
    info['gpu'] = get_gpu_status()

    if info.get('phase') in ('completed', 'done'):
        info['display_step'] = '完成'
        info['display_pct'] = 100.0
    elif info.get('step') and info.get('total_steps'):
        # Calculate overall progress: (epoch-1)*100/3 + step/total*100/3
        ep = info.get('epoch') or 1
        total_ep = info.get('total_epoch') or 3
        ep_progress = (ep - 1) * 100 / total_ep
        ep_step = info['step'] / info['total_steps'] * 100 / total_ep
        info['display_pct'] = round(ep_progress + ep_step, 1)
        if info.get('epoch'):
            info['display_step'] = f"Ep{info['epoch']}/{info['total_epoch']} · {info['step']}/{info['total_steps']}"
        else:
            info['display_step'] = f"{info['step']}/{info['total_steps']}"
    else:
        info['display_step'] = phase if phase != 'no_log' else '等待数据...'
        info['display_pct'] = 0

    info['updated'] = datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%S.000000+00:00Z')
    with open(DATA_FILE, 'w') as f:
        json.dump(info, f, ensure_ascii=False)
    print(json.dumps(info, ensure_ascii=False, indent=2))

if __name__ == '__main__':
    main()