"""
BitFlow MoE v4 部署脚本
将训练好的 int8 权重和推理代码部署到 micro:bit

用法:
  python deploy_moe.py              # 自动查找 micro:bit 并部署
  python deploy_moe.py D:\\          # 指定 micro:bit 盘符
"""
import os
import sys
import shutil
import glob

# 文件路径
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
WEIGHTS_FILE = os.path.join(SCRIPT_DIR, "bitflow_moe_weights.bin")
INFER_CODE = os.path.join(SCRIPT_DIR, "bitflow_moe.py")

# micro:bit 上的文件名
TARGET_WEIGHTS = "weights.bin"
TARGET_MAIN = "main.py"


def find_microbit():
    """自动查找 micro:bit USB 盘符"""
    # Windows: 查找 MICROBIT 卷标
    for letter in "ABCDEFGHIJKLMNOPQRSTUVWXYZ":
        drive = f"{letter}:\\"
        if os.path.exists(drive):
            # 检查是否是 micro:bit
            label_path = os.path.join(drive, "DETAILS.TXT")
            if os.path.exists(label_path):
                try:
                    with open(label_path, "r") as f:
                        content = f.read()
                    if "micro:bit" in content.lower() or "V2" in content or "nRF52" in content:
                        return drive
                except:
                    pass
            # 也检查卷标
            vol_label = os.path.join(drive, "MICROBIT.HTM")
            if os.path.exists(vol_label):
                return drive
    return None


def deploy(target_drive):
    """部署文件到 micro:bit"""
    target_weights = os.path.join(target_drive, TARGET_WEIGHTS)
    target_main = os.path.join(target_drive, TARGET_MAIN)

    # 检查源文件
    if not os.path.exists(WEIGHTS_FILE):
        print(f"[ERROR] 权重文件不存在: {WEIGHTS_FILE}")
        print("  请先运行: python train_bitflow_moe.py")
        return False

    if not os.path.exists(INFER_CODE):
        print(f"[ERROR] 推理代码不存在: {INFER_CODE}")
        return False

    # 权重文件大小
    w_size = os.path.getsize(WEIGHTS_FILE)
    print(f"\n[deploy] 部署到 micro:bit")
    print(f"  盘符: {target_drive}")
    print(f"  权重: {WEIGHTS_FILE} ({w_size/1024:.1f}KB)")
    print(f"  代码: {INFER_CODE}")

    # 拷贝权重
    print(f"\n  [1/2] 拷贝权重 -> {TARGET_WEIGHTS} ...")
    try:
        shutil.copy2(WEIGHTS_FILE, target_weights)
        print(f"       OK ({w_size/1024:.1f}KB)")
    except Exception as e:
        print(f"       FAILED: {e}")
        return False

    # 拷贝代码 (重命名为 main.py)
    print(f"  [2/2] 拷贝代码 -> {TARGET_MAIN} ...")
    try:
        shutil.copy2(INFER_CODE, target_main)
        code_size = os.path.getsize(target_main)
        print(f"       OK ({code_size/1024:.1f}KB)")
    except Exception as e:
        print(f"       FAILED: {e}")
        return False

    # 验证
    print(f"\n[verify] 部署验证:")
    if os.path.exists(target_weights):
        print(f"  {TARGET_WEIGHTS}: {os.path.getsize(target_weights)} bytes OK")
    else:
        print(f"  {TARGET_WEIGHTS}: MISSING!")
        return False

    if os.path.exists(target_main):
        print(f"  {TARGET_MAIN}: {os.path.getsize(target_main)} bytes OK")
    else:
        print(f"  {TARGET_MAIN}: MISSING!")
        return False

    total = os.path.getsize(target_weights) + os.path.getsize(target_main)
    print(f"\n[done] 部署完成! 总计 {total/1024:.1f}KB")
    print(f"\n  micro:bit 将自动重启并运行 main.py")
    print(f"  通过串口 (115200) 发送: GEN:0,1,5|20")
    print(f"  (ID格式: 逗号分隔的字符ID, |后为生成长度)")
    print(f"\n  字符ID对照表见下方:")
    return True


def print_vocab():
    """打印字符ID对照表"""
    VOCAB_CHARS = (
        "猫狗鸟鱼猪牛马鸡虫羊兔鸭鹅"
        "坐跑飞游看吃睡走听来叫爬跳唱画读做"
        "在上里到"
        "的是好了"
        "吗什么怎哪"
        "你我他她们它"
        "蓝绿红黄白"
        "天地水火花草树木日月星云春夏秋冬"
        "大小多高远快"
        "有没会不"
        "人头脚手眼腿"
        "垫肉书子"
        "一三只两"
        "热冷落开"
        "自然动植物质色样颜得都"
        "你好吗谢谢再见早晚安静"
        "很棒爱喜欢乐伤心气"
        "名字叫谁和问答应"
        "聊天话今天说"
        "客用"
        "bit"
        "\n"
    )
    print(f"\n{'='*60}")
    print("字符ID对照表 (用于串口通信)")
    print(f"{'='*60}")
    for i, c in enumerate(VOCAB_CHARS):
        display = "\\n" if c == "\n" else c
        print(f"  {i:3d}: {display}  ", end="")
        if (i + 1) % 8 == 0:
            print()
    print()
    print(f"\n常用短语ID:")
    examples = {
        "你好": "0,1",  # 需要查实际ID
    }
    # 生成常用短语
    phrases = ["你好", "你好吗", "谢谢", "再见", "早安", "晚安",
               "你叫什么名字", "你喜欢什么", "什么会飞", "什么会跑"]
    CH2I = {c: i for i, c in enumerate(VOCAB_CHARS)}
    for p in phrases:
        ids = [str(CH2I[c]) for c in p if c in CH2I]
        print(f"  {p}: {','.join(ids)}")
    print(f"\n  生成命令: GEN:<ID列表>|<生成长度>")
    print(f"  示例:     GEN:{','.join(str(CH2I[c]) for c in '你好' if c in CH2I)}|10")


def main():
    print("=" * 60)
    print("BitFlow MoE v4 部署工具")
    print("=" * 60)

    # 查找 micro:bit
    if len(sys.argv) > 1:
        target = sys.argv[1]
        if not os.path.exists(target):
            print(f"[ERROR] 指定的盘符不存在: {target}")
            return
    else:
        print("\n[search] 正在查找 micro:bit ...")
        target = find_microbit()
        if target is None:
            print("[ERROR] 未找到 micro:bit!")
            print("\n请确保:")
            print("  1. micro:bit 已通过 USB 连接")
            print("  2. micro:bit 已出现在文件管理器中")
            print("  3. 如已知盘符, 可指定: python deploy_moe.py D:\\")
            return

    # 部署
    success = deploy(target)

    if success:
        print_vocab()


if __name__ == "__main__":
    main()