# -*- coding: utf-8 -*-
"""
手动升级 @deepseek-ai/dsh（绕过 npm/npx，本机 npm 写 cacache 卡死）
用法:
  python -S dsh_upgrade.py [目标版本]
  例: python -S dsh_upgrade.py 0.1.1-rc.2
  （不传版本则用 DSH_UPGRADE_VERSION 环境变量或默认值）
"""
import urllib.request, tarfile, io, json, os, shutil, sys, time, glob

# 自动检测最新旧版 dsh 缓存目录（_npx 下含 @deepseek-ai/dsh 的目录，取修改时间最新）
def find_src():
    base = r'C:\Users\maidou\AppData\Local\npm-cache\_npx'
    hits = glob.glob(os.path.join(base, '*', 'node_modules', '@deepseek-ai', 'dsh', 'package.json'))
    if not hits:
        print('错误: 未找到旧版 dsh 缓存目录（_npx 下无 @deepseek-ai/dsh）'); sys.exit(1)
    hits.sort(key=os.path.getmtime, reverse=True)
    pkg = os.path.dirname(hits[0])                    # .../@deepseek-ai/dsh
    return os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(pkg))))
    # 上推4层: node_modules/@deepseek-ai/dsh -> _npx/<hash>

SRC = find_src()
SRC_HASH = os.path.basename(SRC)
DST = os.path.join(r'D:\dsh-install', SRC_HASH)
VERSION = sys.argv[1] if len(sys.argv) > 1 else os.environ.get('DSH_UPGRADE_VERSION', '0.1.1-rc.2')
SCOPE = '@deepseek-ai'
REG = 'https://registry.npmmirror.com'

NM  = os.path.join(DST, 'node_modules')
SCOPE_DIR = os.path.join(NM, SCOPE)

def log(msg):
    print(msg, flush=True)

def download(url, timeout=60):
    req = urllib.request.Request(url, headers={'User-Agent': 'python-dsh-upgrade'})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return r.read()

def pkg_url(name):
    # name: 'dsh-app-boot'
    return f'{REG}/{SCOPE}/{name}/-/{name}-{VERSION}.tgz'

def extract_tarball(data, target):
    tf = tarfile.open(fileobj=io.BytesIO(data), mode='r:gz')
    members = tf.getmembers()
    for m in members:
        if m.isdir():
            continue
        if not m.name.startswith('package/'):
            continue
        rel = m.name[len('package/'):]
        dest = os.path.join(target, rel)
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        src_f = tf.extractfile(m)
        if src_f is None:
            continue
        with open(dest, 'wb') as f:
            f.write(src_f.read())
    tf.close()

def update_pkg(pkg_dir, pkg_name):
    url = pkg_url(pkg_name)
    try:
        data = download(url)
    except Exception as e:
        return False, str(e)
    target = os.path.join(pkg_dir, pkg_name)
    if os.path.exists(target):
        shutil.rmtree(target)
    os.makedirs(target)
    extract_tarball(data, target)
    return True, ''

def main():
    t0 = time.time()
    # ---- 1. 复制旧目录 ----
    if not os.path.exists(DST):
        log(f'[1/4] 复制旧目录 {SRC} -> {DST} (291MB, 约1-2分钟) ...')
        shutil.copytree(SRC, DST)
        log(f'  复制完成 ({time.time()-t0:.0f}s)')
    else:
        log('[1/4] 目标目录已存在，跳过复制')

    # ---- 2. 升级 @deepseek-ai/* 家族 ----
    if not os.path.isdir(SCOPE_DIR):
        log(f'错误: {SCOPE_DIR} 不存在'); sys.exit(1)
    names = sorted(os.listdir(SCOPE_DIR))
    log(f'[2/4] 扫描 {SCOPE} 下 {len(names)} 个包，升级到 {VERSION} ...')
    updated, skipped, failed = [], [], []
    for i, name in enumerate(names, 1):
        ok, err = update_pkg(SCOPE_DIR, name)
        if ok:
            updated.append(name)
            log(f'  [{i}/{len(names)}] UPDATED {name}@{VERSION}')
        else:
            skipped.append(name)
            if i % 20 == 0 or i == len(names):
                log(f'  [{i}/{len(names)}] ... 已跳过 {len(skipped)} 个（无 {VERSION} 版本）')
    log(f'  升级 {len(updated)} 个，跳过 {len(skipped)} 个，失败 {len(failed)} 个')

    # ---- 3. 升级 dsh 本体 ----
    log('[3/4] 升级 @deepseek-ai/dsh 本体 ...')
    ok, err = update_pkg(SCOPE_DIR, 'dsh')
    log(f'  dsh 本体: {"OK" if ok else "FAIL " + err}')

    # ---- 4. 验证 ----
    log('[4/4] 验证版本 ...')
    pj = os.path.join(SCOPE_DIR, 'dsh', 'package.json')
    if os.path.exists(pj):
        with open(pj, encoding='utf-8') as f:
            ver = json.load(f).get('version')
        log(f'  @deepseek-ai/dsh version = {ver}')
    else:
        log('  错误: dsh/package.json 不存在')
    log(f'总耗时 {time.time()-t0:.0f}s')
    log('DONE')

if __name__ == '__main__':
    main()