#!/usr/bin/env python3
"""铸渊训练状态Watchdog - GPU服务器端"""
import os, json, time, re, sys, subprocess
sys.stdout.reconfigure(line_buffering=True)

LOG_FILE = "/root/autodl-tmp/train_mother.log"
CODE_LOG = "/root/autodl-tmp/train_coder.log"
STATUS_FILE = "/root/autodl-tmp/training_status.json"
POLL_INTERVAL = 300

def run(cmd):
    try:
        r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=30)
        return r.stdout.strip(), r.returncode
    except:
        return "", -1

def get_gpu_info():
    out, _ = run("nvidia-smi --query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu --format=csv,noheader,nounits")
    parts = out.split(", ")
    if len(parts) >= 4:
        return {"util_pct": float(parts[0]), "mem_used_mb": int(parts[1]), 
                "mem_total_mb": int(parts[2]), "temp_c": float(parts[3])}
    return None

def parse_mother_status():
    status = {"status": "unknown", "status_label": "⚪ 未知",
              "progress": "N/A", "current_step": 0, "current_loss": None,
              "current_epoch": 0, "elapsed_hours": 0, "has_error": False}
    if not os.path.exists(LOG_FILE):
        status["status"] = "pending"; status["status_label"] = "⚪ 未启动"
        return status
    with open(LOG_FILE, 'r') as f:
        content = f.read()
    if "Traceback" in content:
        status["status"] = "error"; status["status_label"] = "🔴 报错"; status["has_error"] = True
    if "DONE!" in content:
        status["status"] = "completed"; status["status_label"] = "🟢 已完成"
        return status
    if "[5/5] Starting training!" in content:
        status["status"] = "training"; status["status_label"] = "🟢 训练中"
        steps = re.findall(r"global_step[\s:=]+(\d+)|Step:\s*(\d+)", content)
        losses = re.findall(r"loss['\s:=]+([\d.]+)", content)
        if steps: status["current_step"] = int([s for pair in steps for s in pair if s][-1])
        if losses: status["current_loss"] = float(losses[-1])
    elif "Tokenize:" in content:
        matches = re.findall(r"Tokenize:\s*\d+%\|.*?\|\s*(\d+)/(\d+)", content)
        if matches:
            done, total = matches[-1]
            status["progress"] = f"{done} / {total} 条"
            pct = int(done) / int(total) * 100
            status["status_label"] = f"🟡 分词中 ({pct:.0f}%)"
            status["status"] = "tokenizing"
    return status

def parse_coder_status():
    if not os.path.exists(CODE_LOG):
        return {"status": "pending", "status_label": "⚪ 未启动"}
    with open(CODE_LOG, 'r') as f:
        content = f.read()
    if "Traceback" in content:
        return {"status": "error", "status_label": "🔴 报错", "has_error": True}
    if "DONE!" in content:
        return {"status": "completed", "status_label": "🟢 已完成"}
    if "[5/5] Starting training!" in content:
        return {"status": "training", "status_label": "🟢 训练中"}
    return {"status": "starting", "status_label": "🟡 启动中"}

def check_output():
    mother_out = "/root/autodl-tmp/output/qwen25-7b-sft/final"
    coder_out = "/root/autodl-tmp/output/qwen25-coder-7b-sft/final"
    return {
        "mother_exists": os.path.isdir(mother_out) and bool(os.listdir(mother_out)),
        "coder_exists": os.path.isdir(coder_out) and bool(os.listdir(coder_out))
    }

def main():
    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] Watchdog started"); sys.stdout.flush()
    while True:
        try:
            gpu = get_gpu_info()
            mother = parse_mother_status()
            coder = parse_coder_status()
            output = check_output()
            status = {
                "last_updated": time.strftime("%Y-%m-%dT%H:%M:%S+08:00"),
                "mother_model": {**mother, "gpu": gpu},
                "code_model": coder,
                "output": output,
                "alerts": []
            }
            if mother.get("has_error"):
                status["alerts"].append({"level": "error", "message": "❌ 母模型报错！找铸渊"})
            if coder.get("has_error"):
                status["alerts"].append({"level": "error", "message": "❌ 代码模型报错！找铸渊"})
            with open(STATUS_FILE, 'w') as f:
                json.dump(status, f, indent=2, ensure_ascii=False)
            gpu_str = f"GPU:{gpu['mem_used_mb']}/{gpu['mem_total_mb']}MB {gpu['temp_c']}C" if gpu else "N/A"
            print(f"[{time.strftime('%H:%M:%S')}] {mother['status_label']} | {gpu_str}")
            sys.stdout.flush()
        except Exception as e:
            print(f"Error: {e}"); sys.stdout.flush()
        time.sleep(POLL_INTERVAL)

if __name__ == "__main__":
    main()