"""
大文件上传脚本 - 通过 JumpServer plink PTY 使用 base64 通道上传
使用 4000 字符的大块，比 jms_ops.py 的 500 字符快 8 倍
"""
import sys, os, base64, hashlib, time

# 导入 jms_ops 模块
SKILL_DIR = r"C:\Users\simon\.workbuddy\skills\jumpserver-ops\scripts"
sys.path.insert(0, SKILL_DIR)
import jms_ops

def upload_large_file(local_path, remote_path, chunk_size=4000):
    """通过 plink PTY 上传大文件，使用大块 base64 编码"""
    local_data = open(local_path, 'rb').read()
    local_md5 = hashlib.md5(local_data).hexdigest()
    b64_data = base64.b64encode(local_data).decode('ascii')

    total_chunks = (len(b64_data) + chunk_size - 1) // chunk_size
    print(f"📤 上传: {local_path} → {remote_path}")
    print(f"   原始: {len(local_data)} bytes, base64: {len(b64_data)} chars")
    print(f"   分 {total_chunks} 块发送 (每块 {chunk_size} chars)")

    # 获取 token + 启动会话
    tokens = jms_ops.get_connection_tokens(1)
    if not tokens:
        print("❌ 获取 token 失败")
        return False

    token_id, token_secret = tokens[0]
    session = jms_ops.PlinkSession(f"JMS-{token_id}", token_secret)
    if not session.connect():
        print("❌ 会话启动失败")
        return False

    try:
        # 1. 清空目标文件
        session.run_command(f'> {remote_path}', timeout=5)

        # 2. 逐块追加 (大块)
        start_time = time.time()
        for i in range(0, len(b64_data), chunk_size):
            chunk = b64_data[i:i+chunk_size]
            chunk_num = i // chunk_size + 1
            cmd = f"echo '{chunk}' | base64 -d >> {remote_path}"
            r = session.run_command(cmd, timeout=15)
            if not r["success"]:
                print(f"  ❌ 块 {chunk_num}/{total_chunks} 发送失败")
                return False
            if chunk_num % 50 == 0 or chunk_num == total_chunks:
                elapsed = time.time() - start_time
                pct = chunk_num / total_chunks * 100
                print(f"  📦 已发送 {chunk_num}/{total_chunks} 块 ({pct:.0f}%) - {elapsed:.1f}s")

        # 3. 验证大小
        r = session.run_command(f'wc -c < {remote_path}', timeout=5)
        if r["success"]:
            remote_size = int(r["output"].strip()) if r["output"].strip().isdigit() else -1
            if remote_size == len(local_data):
                elapsed = time.time() - start_time
                print(f"  ✅ 上传成功! 大小匹配 ({remote_size} bytes), 耗时 {elapsed:.1f}s")
                # 4. MD5 验证
                r2 = session.run_command(f'md5sum {remote_path}', timeout=5)
                if r2["success"]:
                    remote_md5 = r2["output"].split()[0]
                    if remote_md5 == local_md5:
                        print(f"  ✅ MD5 匹配! 文件完整")
                    else:
                        print(f"  ⚠️ MD5 不匹配 (本地 {local_md5[:12]}, 远程 {remote_md5[:12]})")
                return True
            else:
                print(f"  ❌ 大小不匹配 (本地 {len(local_data)}, 远程 {remote_size})")
                return False
        else:
            print("  ⚠️ 无法验证远程文件大小")
            return False
    finally:
        session.close()

if __name__ == '__main__':
    local = r"D:\资料\03-项目开发\wecom_it_smart_desk\tmp-agent-dist.tar.gz"
    remote = "/tmp/agent-dist.tar.gz"
    success = upload_large_file(local, remote, chunk_size=4000)
    if success:
        print("\n✅ 上传完成，可以在服务器上解压了")
    else:
        print("\n❌ 上传失败")
        sys.exit(1)