"""流水线状态文件 ml/runs/pipeline.json（供训练监控网页显示"现在在哪一步、每步用了多久"）。
pipeline.sh 每个阶段开始/结束时自动调用，不需要手动运行：
  python -m ml.pstate reset <阶段1,阶段2,...>        # `all` 开始时：清空并写入计划
  python -m ml.pstate begin <阶段> [键=值 ...]        # 阶段开始（值可为数字），如 target_tasks=50000
  python -m ml.pstate set <阶段> 键=值 ...            # 阶段中途更新（如补量后的目标条数）
  python -m ml.pstate end <阶段> <返回码>
2026-10-18 新增（用户要求：从训练开始到结束都能在监控面板看到进度）。"""
import json, os, sys, time

P = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'ml', 'runs', 'pipeline.json')


def load():
    try:
        with open(P, encoding='utf-8') as f: return json.load(f)
    except Exception:
        return {'started': time.time(), 'plan': None, 'stages': {}, 'current': None}


def save(d):
    os.makedirs(os.path.dirname(P), exist_ok=True)
    tmp = P + '.tmp'
    with open(tmp, 'w', encoding='utf-8') as f: json.dump(d, f, ensure_ascii=False, indent=1)
    os.replace(tmp, P)


def kv(args):
    out = {}
    for a in args:
        k, _, v = a.partition('=')
        try: v = float(v) if '.' in v else int(v)
        except ValueError: pass
        out[k] = v
    return out


def main():
    cmd, *rest = sys.argv[1:] or ['?']
    d = load(); now = time.time()
    if cmd == 'reset':
        plan = rest[0].split(',') if rest else []
        # 从中间某步续跑时保留计划之外（之前已跑完）各步的记录，计划内的清空重来
        d = {'started': now, 'plan': plan, 'stages': {k: v for k, v in d.get('stages', {}).items() if k not in plan}, 'current': None}
    elif cmd == 'begin':
        st = rest[0]
        d['stages'][st] = {'start': now, 'end': None, 'rc': None, 'info': kv(rest[1:])}
        d['current'] = st
    elif cmd == 'set':
        s = d['stages'].setdefault(rest[0], {'start': now, 'end': None, 'rc': None, 'info': {}})
        s['info'].update(kv(rest[1:]))
    elif cmd == 'end':
        st, rc = rest[0], int(rest[1]) if len(rest) > 1 else 0
        s = d['stages'].setdefault(st, {'start': now, 'info': {}})
        s['end'], s['rc'] = now, rc
        if d.get('current') == st: d['current'] = None
    else:
        print(__doc__); return
    save(d)


if __name__ == '__main__':
    main()