#!/usr/bin/env python3
"""
蒸馏训练监控面板 — 在你本地终端运行
依赖：sshpass (brew install sshpass)

用法：
  export ZY_SSH_PASS="你的SSH密码"
  python3 scripts/monitor_distill.py

或者放在 ~/.zshrc 里就不用每次export了
"""

import subprocess, re, sys, time, shutil, os

SSH_HOST = os.environ.get("ZY_SSH_HOST", "connect.westd.seetacloud.com")
SSH_PORT = os.environ.get("ZY_SSH_PORT", "23647")
SSH_USER = os.environ.get("ZY_SSH_USER", "root")
# 从环境变量读取，不硬编码
SSH_PASS = os.environ.get("ZY_SSH_PASS", "")
LOG_FILE = "/root/autodl-tmp/distill_mother.log"
TRAINING_LOG = "/root/autodl-tmp/train_coder.log"

def ssh(cmd):
    try:
        r = subprocess.run(
            ["sshpass", "-p", SSH_PASS, "ssh", "-o", "StrictHostKeyChecking=accept-new",
             "-o", "ConnectTimeout=5", "-p", SSH_PORT,
             f"{SSH_USER}@{SSH_HOST}", cmd],
            capture_output=True, text=True, timeout=10
        )
        return r.stdout, r.stderr
    except:
        return "", "timeout"

def get_bar(percent, width=20):
    filled = int(percent / 100 * width)
    return "\u2588" * filled + "\u2591" * (width - filled)

def main():
    print("=" * 60)
    print("  铸渊蒸馏监控面板")
    print("  Qwen2.5-7B \u2192 Qwen2.5-1.5B (霜砚模板)")
    print("=" * 60)
    
    while True:
        stdout, stderr = ssh(f"tail -30 {LOG_FILE} 2>/dev/null; echo '---GPU---'; nvidia-smi --query-gpu=memory.used,memory.total,utilization.gpu,temperature.gpu --format=csv,noheader 2>/dev/null; echo '---DISK---'; df -h /root/autodl-tmp 2>/dev/null | tail -1; echo '---PROC---'; ps aux | grep distill_mother | grep -v grep | wc -l")
        
        lines = stdout.strip().split('\n')
        
        log_lines = []
        gpu_info = ""
        disk_info = ""
        proc_count = "0"
        
        in_gpu = False
        in_disk = False
        in_proc = False
        
        for line in lines:
            if line == '---GPU---':
                in_gpu = True; in_disk = False; in_proc = False; continue
            elif line == '---DISK---':
                in_gpu = False; in_disk = True; in_proc = False; continue
            elif line == '---PROC---':
                in_gpu = False; in_disk = False; in_proc = True; continue
            
            if in_gpu:
                gpu_info = line
            elif in_disk:
                disk_info = line
            elif in_proc:
                proc_count = line.strip()
            else:
                log_lines.append(line)
        
        progress = 0
        loss_str = ""
        epoch_str = ""
        step_str = ""
        eta = ""
        
        for line in log_lines:
            m = re.search(r'(\d+)%\|', line)
            if m:
                progress = int(m.group(1))
            
            m = re.search(r'\|[^\|]*\|\s*(\d+)/(\d+)', line)
            if m:
                step_str = f"{m.group(1)}/{m.group(2)}"
            
            m = re.search(r"'loss':\s*'?([\d.]+)'?", line)
            if m:
                loss_str = m.group(1)
            
            m = re.search(r"'epoch':\s*'?([\d.]+)'?", line)
            if m:
                epoch_str = m.group(1)
            
            m = re.search(r'<(\d+:\d+)', line)
            if m:
                eta = m.group(1)
        
        subprocess.run(["clear"])
        term_width = shutil.get_terminal_size().columns
        
        print("=" * term_width)
        title = "  \U0001f504 蒸馏训练监控 "
        print(f"{title:=^{term_width}}")
        print("=" * term_width)
        
        bar = get_bar(progress)
        print(f"\n  进度: [{bar}] {progress}%")
        if step_str:
            print(f"  步数: {step_str}  Epoch: {epoch_str or '?'}")
        if loss_str:
            print(f"  Loss: {loss_str}")
        if eta:
            print(f"  剩余: ~{eta}")
        
        print()
        print("-" * term_width)
        
        if gpu_info:
            parts = gpu_info.split(', ')
            if len(parts) >= 2:
                used, total = parts[0].strip(), parts[1].strip()
                try:
                    mem_pct = int(used.split()[0]) / int(total.split()[0]) * 100
                    mem_bar = get_bar(mem_pct)
                    print(f"  GPU显存: [{mem_bar}] {used}/{total} ({mem_pct:.0f}%)")
                except:
                    print(f"  GPU显存: {used}/{total}")
                if len(parts) >= 3:
                    print(f"  利用率: {parts[3] if len(parts)>=4 else parts[2]}")
        
        if disk_info:
            parts = disk_info.split()
            if len(parts) >= 5:
                used_pct = parts[4].rstrip('%')
                try:
                    pct = int(used_pct)
                    disk_bar = get_bar(pct)
                    print(f"  数据盘: [{disk_bar}] {parts[2]}/{parts[1]} ({used_pct}%)")
                except:
                    print(f"  数据盘: {parts[2]}/{parts[1]} ({used_pct}%)")
        
        if proc_count == "0":
            print(f"\n  \u26a0\ufe0f  蒸馏进程未运行！可能已完成或出错")
        else:
            print(f"\n  \u2705 蒸馏训练进行中")
        
        print()
        print("-" * term_width)
        print("  最近日志:")
        recent = [l for l in log_lines if l.strip() and not l.startswith('\r')][-5:]
        for l in recent[-3:]:
            cleaned = l.replace('\r', '').strip()
            if cleaned:
                print(f"    {cleaned[:100]}")
        
        print()
        print(f"  刷新: 5秒 | {time.strftime('%H:%M:%S')}")
        print("  Ctrl+C 退出")
        print("=" * term_width)
        
        time.sleep(5)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n  监控已退出")
        sys.exit(0)