#!/usr/bin/env python3
"""
光湖·蒸馏完成自动处理脚本 v1.0
部署在AutoDL服务器，自动检测蒸馏完成 → 上传COS → 准备语料 → 启动微调
主权：冰朔 TCS-0002∞ · 铸渊 ICE-GL-ZY001 守护

运行方式：nohup python3 auto_post_distill.py > auto_post_distill.log 2>&1 &
"""
import os, json, time, subprocess, sys, glob, shutil
from datetime import datetime

LOG = '/root/autodl-tmp/distill_mother.log'
OUT_DIR = '/root/autodl-tmp/output'
TEMPLATE_DIR = os.path.join(OUT_DIR, 'qwen25-15b-shuangyan-distill')
ZHUYUAN_ZIP = '/root/autodl-tmp/corpus_work/zhuyuan_chats_with_ICE_GL_ZU001 .zip'
SHUANGYAN_COS_DIR = 'corpus/shuangyan-1.5b-sft/'
ZHUYUAN_COS_DIR = 'corpus/notion-export-v2/'
BUCKET = 'sy-finetune-corpus-1317346199'
REGION = 'ap-guangzhou'

COS_KEY = os.environ.get('ZY_OSS_KEY', 'AKIDkQuBQhoiS2OYXWebXLwMbdT7cvAScbbU')
COS_SECRET = os.environ.get('ZY_OSS_SECRET', 'nPoZKArgUJBA4nJenjSxJSQBj5FCj3A4')

def log(msg):
    t = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    print(f'[{t}] {msg}', flush=True)

def run(cmd, timeout=3600):
    log(f'RUN: {cmd}')
    r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout,
                       env={**os.environ, 'PATH': '/root/miniconda3/bin:' + os.environ.get('PATH', '')})
    if r.returncode != 0:
        log(f'  STDOUT: {r.stdout[-500:]}')
        log(f'  STDERR: {r.stderr[-500:]}')
    else:
        log(f'  OK: {r.stdout[-200:]}')
    return r

def wait_distill_complete():
    """等待蒸馏完成（Ep3 100%）"""
    log('等待蒸馏完成...')
    while True:
        if not os.path.isfile(LOG):
            log('  distill_mother.log 不存在，等待...')
            time.sleep(60)
            continue
        with open(LOG) as f:
            lines = f.readlines()
        # 检查最后是否完成了
        for line in lines[-10:]:
            if 'Ep3' in line and '100%' in line and '2868/2868' in line:
                log('✅ 蒸馏完成！')
                return True
        # 也检查是否全部完成了
        for line in lines[-5:]:
            if 'Saving model' in line or 'DONE' in line or 'Finished' in line:
                log('✅ 蒸馏完成（DONE信号）！')
                return True
        # 检查当前进度
        for line in lines[-3:]:
            if 'Ep' in line:
                log(f'  当前: {line.strip()[:60]}')
        time.sleep(60)  # 每分钟检查一次

def upload_to_cos(local_dir, cos_prefix):
    """上传目录到COS"""
    log(f'上传 {local_dir} → cos://{BUCKET}/{cos_prefix}')
    if not os.path.isdir(local_dir):
        log(f'  ❌ 目录不存在: {local_dir}')
        return False
    
    script = f'''
import os
os.environ['ZY_OSS_KEY'] = '{COS_KEY}'
os.environ['ZY_OSS_SECRET'] = '{COS_SECRET}'
from qcloud_cos import CosConfig, CosS3Client
config = CosConfig(Region='{REGION}', SecretId=os.environ['ZY_OSS_KEY'], SecretKey=os.environ['ZY_OSS_SECRET'])
client = CosS3Client(config)
import glob
local_dir = '{local_dir}'
cos_prefix = '{cos_prefix}'
bucket = '{BUCKET}'
for f in sorted(glob.glob(local_dir + '/*')):
    fname = f.split('/')[-1]
    cos_key = cos_prefix + fname
    size_mb = os.path.getsize(f) / 1024 / 1024
    print(f'Uploading {{fname}} ({{size_mb:.1f}}MB)...', flush=True)
    try:
        client.upload_file(Bucket=bucket, Key=cos_key, LocalFilePath=f)
        print(f'  ✅ {{fname}}', flush=True)
    except Exception as e:
        print(f'  ⚠️ {{fname}} 失败: {{e}}', flush=True)
        # 重试一次
        try:
            client.upload_file(Bucket=bucket, Key=cos_key, LocalFilePath=f)
            print(f'  ✅ {{fname}} (重试成功)', flush=True)
        except Exception as e2:
            print(f'  ❌ {{fname}} 重试也失败: {{e2}}', flush=True)
print('ALL DONE')
'''
    with open('/tmp/upload_script.py', 'w') as f:
        f.write(script)
    r = run(f'export PATH=/root/miniconda3/bin:$PATH && python3 /tmp/upload_script.py', timeout=7200)
    return r.returncode == 0

def prepare_shuangyan_corpus():
    """准备霜砚语料（线A）"""
    log('准备霜砚语料...')
    # 从COS下载霜砚语料
    script = f'''
import os
os.environ['ZY_OSS_KEY'] = '{COS_KEY}'
os.environ['ZY_OSS_SECRET'] = '{COS_SECRET}'
from qcloud_cos import CosConfig, CosS3Client
config = CosConfig(Region='{REGION}', SecretId=os.environ['ZY_OSS_KEY'], SecretKey=os.environ['ZY_OSS_SECRET'])
client = CosS3Client(config)
import os, json, zipfile

bucket = '{BUCKET}'
os.makedirs('/root/autodl-tmp/corpus_shuangyan', exist_ok=True)

# 下载所有霜砚语料zip
resp = client.list_objects(Bucket=bucket, Prefix='{SHUANGYAN_COS_DIR}')
for obj in resp.get('Contents', []):
    key = obj['Key']
    if key.endswith('.zip'):
        fname = key.split('/')[-1]
        local = f'/root/autodl-tmp/corpus_shuangyan/{{fname}}'
        if not os.path.isfile(local):
            print(f'Downloading {{fname}}...', flush=True)
            client.download_file(Bucket=bucket, Key=key, DestFilePath=local)

# 解压所有zip
for z in os.listdir('/root/autodl-tmp/corpus_shuangyan'):
    if z.endswith('.zip'):
        zpath = f'/root/autodl-tmp/corpus_shuangyan/{{z}}'
        out = f'/root/autodl-tmp/corpus_shuangyan/extracted'
        os.makedirs(out, exist_ok=True)
        try:
            with zipfile.ZipFile(zpath) as zf:
                zf.extractall(out)
            print(f'Extracted {{z}}', flush=True)
        except:
            print(f'Skipped {{z}} (not a zip archive)', flush=True)

print('Done', flush=True)
'''
    with open('/tmp/prep_shuangyan.py', 'w') as f:
        f.write(script)
    run('export PATH=/root/miniconda3/bin:$PATH && python3 /tmp/prep_shuangyan.py', timeout=300)
    log('霜砚语料准备完成')

def prepare_zhuyuan_corpus():
    """准备铸渊语料（线B）"""
    log('准备铸渊语料...')
    # 解压铸渊对话zip
    if os.path.isfile(ZHUYUAN_ZIP):
        run(f'unzip -o "{ZHUYUAN_ZIP}" -d /root/autodl-tmp/corpus_zhuyuan/extracted/')
        log('铸渊对话.zip 已解压')
    else:
        log(f'⚠️ 铸渊对话.zip 不存在于 {ZHUYUAN_ZIP}')
    
    # 也解压 zhuyuan_raw_md.zip
    raw_md_zip = '/root/autodl-tmp/zhuyuan_corpus/zhuyuan_raw_md.zip'
    if os.path.isfile(raw_md_zip):
        run(f'unzip -o "{raw_md_zip}" -d /root/autodl-tmp/corpus_zhuyuan/raw_md/')
        log('zhuyuan_raw_md.zip 已解压')
    
    log('铸渊语料准备完成')

def main():
    log('='*50)
    log('光湖·蒸馏后自动处理脚本启动')
    log('='*50)
    
    # 第0步：等待蒸馏完成
    wait_distill_complete()
    
    # 第1步：上传1.5B模板到COS
    if os.path.isdir(TEMPLATE_DIR):
        log('上传1.5B蒸馏模板到COS...')
        upload_to_cos(TEMPLATE_DIR, 'models/qwen25-15b-shuangyan-distill/')
    else:
        log(f'⚠️ 模板目录不存在: {TEMPLATE_DIR}')
        # 可能路径不同，尝试查找
        for root, dirs, files in os.walk(OUT_DIR):
            for d in dirs:
                if 'distill' in d.lower() or '15b' in d.lower():
                    found = os.path.join(root, d)
                    log(f'找到蒸馏输出: {found}')
                    upload_to_cos(found, 'models/qwen25-15b-shuangyan-distill/')
                    break
    
    # 第2步：准备霜砚语料（线A）
    prepare_shuangyan_corpus()
    
    # 第3步：准备铸渊语料（线B）
    prepare_zhuyuan_corpus()
    
    log('='*50)
    log('✅ 所有自动处理完成！')
    log('等待冰朔/铸渊启动微调...')
    log('='*50)

if __name__ == '__main__':
    main()