# -*- coding: utf-8 -*-
"""
红外板 — 发送协议穷举测试
========================
从接收已知: 客户码=3BC4, 方向键=0D/10/15/12
逐个尝试常见红外板发送格式，找出正确的那一个。

用法:
    python ir_bruteforce.py -p COM8            # 自动试探所有协议
    python ir_bruteforce.py -p COM8 -k 0D      # 指定键值测试
    python ir_bruteforce.py -p COM8 --list     # 列出所有候选协议

每次发送后观察投影仪是否有反应，按 y/n 确认，脚本记录有效协议。
"""

import sys
import time
import argparse
import itertools
import serial
import serial.tools.list_ports


BAUDRATE = 9600

# 已知参数
CUSTOMER_CODE = bytes([0x3B, 0xC4])  # 客户码
KEY_CODES = {
    "up":    0x0D,  # 上
    "left":  0x10,  # 左
    "down":  0x15,  # 下
    "right": 0x12,  # 右
}


def build_frame(proto_name: str, key: int) -> bytes | None:
    """
    根据协议名构建发送帧。

    常见红外板发送协议:
      - raw3:        直接 3 字节 [客户码H, 客户码L, 键值]
      - nec_std:     NEC 标准帧 [地址, ~地址, 命令, ~命令]
      - nec_3B:      地址 0x3B (单字节) 的 NEC [0x3B, 0xC4, key, ~key]
      - a5_prefix:   A5 前缀协议 [A5, len, 数据..., checksum]
      - aa_prefix:   AA 前缀协议 [AA, 数据..., sum]
      - text_crlf:   文本协议 "SEND xxxxxx\\r\\n"
      - text_no_cr:  文本协议 "3BC40D"
      - at_cmd:      AT 指令 "AT+IR=xxxxxx\\r\\n"
      - nec_full:    NEC 完整帧 [addrL, addrH, cmd, ~cmd] (带反码)
    """

    hi, lo = CUSTOMER_CODE[0], CUSTOMER_CODE[1]

    # ── 原始字节类 ──
    if proto_name == "raw3":
        return bytes([hi, lo, key])

    if proto_name == "raw3_reversed":
        return bytes([key, lo, hi])

    if proto_name == "raw3_with_cr":
        return bytes([hi, lo, key, 0x0D])

    if proto_name == "raw3_with_crlf":
        return bytes([hi, lo, key, 0x0D, 0x0A])

    # ── NEC 标准变体 ──
    if proto_name == "nec_3bytes":
        # 最简 NEC: 地址 + 命令 (不带反码)
        return bytes([hi, lo, key])

    if proto_name == "nec_addr_inv":
        # NEC: [地址, ~地址, 命令, ~命令]
        return bytes([hi, hi ^ 0xFF, key, key ^ 0xFF])

    if proto_name == "nec_addr_inv_both":
        # NEC 有时用两字节地址: [addrL, addrH, cmd, ~cmd]
        return bytes([lo, hi, key, key ^ 0xFF])

    if proto_name == "nec_no_invert":
        # 只发地址+命令，不做反码
        return bytes([lo, key])

    if proto_name == "nec_lo_cmd":
        # [0x3B, key] — 只有客户码低字节+键值
        return bytes([lo, key])

    if proto_name == "nec_hi_cmd":
        # [0x3B, 0xC4, key] — raw3 即此格式
        return bytes([hi, key])

    # ── 带前缀/帧头 ──
    if proto_name == "a5_f5":
        # 常见: [A5, len, data..., sum]
        data = bytes([hi, lo, key])
        frame = bytes([0xA5, len(data)]) + data
        csum = sum(frame) & 0xFF
        return frame + bytes([csum])

    if proto_name == "a5_f5_nosum":
        data = bytes([hi, lo, key])
        return bytes([0xA5, len(data)]) + data

    if proto_name == "aa_55":
        data = bytes([hi, lo, key])
        csum = sum(data) & 0xFF
        return bytes([0xAA, 0x55]) + data + bytes([csum])

    if proto_name == "aa_55_nosum":
        return bytes([0xAA, 0x55, hi, lo, key])

    if proto_name == "ff_prefix":
        return bytes([0xFF, 0xFF, hi, lo, key])

    if proto_name == "fe_prefix":
        return bytes([0xFE, hi, lo, key, 0xFD])

    if proto_name == "a0_prefix":
        return bytes([0xA0, hi, lo, key, 0x0A])

    # ── 文本协议类 ──
    hex_str = f"{hi:02X}{lo:02X}{key:02X}"

    if proto_name == "text_hex":
        return hex_str.encode("ascii")

    if proto_name == "text_hex_cr":
        return (hex_str + "\r").encode("ascii")

    if proto_name == "text_hex_crlf":
        return (hex_str + "\r\n").encode("ascii")

    if proto_name == "text_hex_nl":
        return (hex_str + "\n").encode("ascii")

    if proto_name == "text_send":
        return f"SEND {hex_str}".encode("ascii")

    if proto_name == "text_send_crlf":
        return f"IR_SEND={hex_str}\r\n".encode("ascii")

    if proto_name == "text_at_crlf":
        return f"AT+IR={hex_str}\r\n".encode("ascii")

    if proto_name == "text_at_send_crlf":
        return f"AT+SEND={hex_str}\r\n".encode("ascii")

    # ── 带空格分隔的十六进制文本 ──
    if proto_name == "text_spaced_crlf":
        return f"{hi:02X} {lo:02X} {key:02X}\r\n".encode("ascii")

    if proto_name == "text_spaced_nl":
        return f"{hi:02X} {lo:02X} {key:02X}\n".encode("ascii")

    # ── 先发命令再发数据 ──
    if proto_name == "cmd_tx_data":
        # 先发 TX 命令，延时 50ms，再发数据
        return None  # 特殊处理

    if proto_name == "cmd_0x01_prefix":
        return bytes([0x01, hi, lo, key])

    if proto_name == "cmd_0x02_prefix":
        return bytes([0x02, hi, lo, key])

    if proto_name == "cmd_0x03_prefix":
        return bytes([0x03, hi, lo, key])

    return None


# ═══════════════════════════════════════════════════════════════
# 所有候选协议 (按可能性排序)
# ═══════════════════════════════════════════════════════════════

PROTOCOLS = [
    # ── 最可能的: 二进制直发 ──
    ("raw3",              "二进制3字节 [3B C4 XX]"),
    ("raw3_reversed",     "二进制3字节反转 [XX C4 3B]"),
    ("raw3_with_cr",      "二进制3字节 + CR [3B C4 XX 0D]"),
    ("raw3_with_crlf",    "二进制3字节 + CRLF [3B C4 XX 0D 0A]"),

    # ── NEC 变体 ──
    ("nec_3bytes",        "NEC简化 3字节 [3B C4 XX]"),
    ("nec_addr_inv",      "NEC标准 [addr, ~addr, cmd, ~cmd]"),
    ("nec_addr_inv_both", "NEC双字节地址 [C4 3B XX ~XX]"),
    ("nec_lo_cmd",        "NEC短帧 [C4 XX]"),
    ("nec_hi_cmd",        "NEC短帧 [3B XX]"),
    ("nec_no_invert",     "NEC无反码 [C4 XX]"),

    # ── 带帧头 ──
    ("a5_f5",             "A5帧头+校验 [A5 03 3B C4 XX SUM]"),
    ("a5_f5_nosum",       "A5帧头无校验 [A5 03 3B C4 XX]"),
    ("aa_55",             "AA55帧头 [AA 55 3B C4 XX SUM]"),
    ("aa_55_nosum",       "AA55帧头无校验 [AA 55 3B C4 XX]"),
    ("ff_prefix",         "FF前缀 [FF FF 3B C4 XX]"),
    ("fe_prefix",         "FE帧 [FE 3B C4 XX FD]"),
    ("a0_prefix",         "A0前缀 [A0 3B C4 XX 0A]"),
    ("cmd_0x01_prefix",   "01前缀 [01 3B C4 XX]"),
    ("cmd_0x02_prefix",   "02前缀 [02 3B C4 XX]"),
    ("cmd_0x03_prefix",   "03前缀 [03 3B C4 XX]"),

    # ── 文本协议 ──
    ("text_hex",          "文本HEX \"3BC40D\""),
    ("text_hex_cr",       "文本HEX+CR"),
    ("text_hex_crlf",     "文本HEX+CRLF"),
    ("text_hex_nl",       "文本HEX+LF"),
    ("text_spaced_crlf",  "文本空格HEX \"3B C4 0D\\r\\n\""),
    ("text_spaced_nl",    "文本空格HEX \"3B C4 0D\\n\""),
    ("text_send",         "文本 \"SEND 3BC40D\""),
    ("text_send_crlf",    "文本 \"IR_SEND=3BC40D\\r\\n\""),
    ("text_at_crlf",      "AT指令 \"AT+IR=3BC40D\\r\\n\""),
    ("text_at_send_crlf", "AT指令 \"AT+SEND=3BC40D\\r\\n\""),
]


def scan_usb_ttl():
    """自动找 USB-to-TTL 端口"""
    ports = serial.tools.list_ports.comports()
    keywords = ["ch343", "ch340", "ch341", "usb-enhanced-serial",
                "usb serial", "ft232", "cp210", "pl2303", "wch"]
    for p in ports:
        combined = f"{p.description or ''} {p.hwid or ''}".lower()
        if any(kw in combined for kw in keywords):
            return p.device
    all_p = list(ports)
    return all_p[0].device if all_p else None


def brute_force(port: str, key_hex: str = "0D", start_from: int = 0):
    """
    逐个尝试所有协议。
    每次发送后等待用户确认 (y=有效/n=无效/q=退出)
    """
    key = int(key_hex, 16)
    hi, lo = CUSTOMER_CODE[0], CUSTOMER_CODE[1]
    hex_3b = f"{hi:02X}{lo:02X}{key:02X}"

    print(f"串口: {port} @ {BAUDRATE}bps")
    print(f"客户码: {hi:02X} {lo:02X}")
    print(f"测试键值: {key:02X} ({hex_3b})")
    print(f"候选协议: {len(PROTOCOLS)} 个\n")

    ser = serial.Serial(port, baudrate=BAUDRATE, timeout=0.3)
    print("=" * 65)
    print("  发送后观察投影仪是否响应，输入 y/n/q")
    print("=" * 65)

    working = []
    total = len(PROTOCOLS)
    skipped = 0

    for i, (name, desc) in enumerate(PROTOCOLS):
        if i < start_from:
            skipped += 1
            continue

        frame = build_frame(name, key)

        # 特殊协议: 先发送 "TX" 命令再发数据
        if name == "cmd_tx_data":
            try:
                ser.write(b"TX\r\n")
                ser.flush()
                time.sleep(0.05)
                ser.write(bytes([hi, lo, key]))
                ser.flush()
            except Exception as e:
                print(f"  [{i+1}/{total}] {name}: 发送失败 - {e}")
                continue
            print(f"\n[{i+1}/{total}] {name}")
            print(f"  描述: {desc}")
            print(f"  发送: TX\\r\\n → [3B C4 {key:02X}]")

        elif frame is not None:
            try:
                ser.reset_output_buffer()
                ser.write(frame)
                ser.flush()
            except Exception as e:
                print(f"  [{i+1}/{total}] {name}: 发送失败 - {e}")
                continue

            print(f"\n[{i+1}/{total}] {name}")
            print(f"  描述: {desc}")
            print(f"  HEX:  {frame.hex(' ').upper()}")
            if all(32 <= b < 127 for b in frame):
                print(f"  ASCII: {frame.decode('ascii')}")

        else:
            continue

        # 用户确认
        while True:
            resp = input("  有效? [y/n/q/s(跳过)]: ").strip().lower()
            if resp == 'y':
                working.append((name, desc, frame.hex(" ").upper() if frame else "TX+data"))
                print(f"  ✓ 已记录")
                break
            elif resp == 'n':
                break
            elif resp == 'q':
                print(f"\n已测试 {i+1-skipped}/{total}")
                ser.close()
                return working
            elif resp == 's':
                print(f"  跳过")
                break
            else:
                print("  请输入 y(有效) / n(无效) / q(退出)")

        time.sleep(0.1)

    ser.close()

    print(f"\n测试完成: {i+1-skipped}/{total}")
    return working


def main():
    parser = argparse.ArgumentParser(
        description="红外板发送协议 — 穷举测试",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("-p", "--port", default=None, help="串口 (如 COM8)")
    parser.add_argument("-k", "--key", default="0D",
                        help=f"测试键值 hex (默认 0D=上). 已知: {', '.join(f'{k}={v:02X}' for k,v in KEY_CODES.items())}")
    parser.add_argument("-s", "--start", type=int, default=0,
                        help="从第 N 个协议开始 (断点续测)")
    parser.add_argument("--list", action="store_true", help="列出所有候选协议")
    args = parser.parse_args()

    if args.list:
        print("候选发送协议:")
        for i, (name, desc) in enumerate(PROTOCOLS):
            print(f"  [{i+1:>2}] {name:<22} {desc}")
        return

    port = args.port or scan_usb_ttl()
    if not port:
        print("✗ 未找到串口。用 -p 指定")
        return

    results = brute_force(port, args.key, start_from=args.start)

    if results:
        print("\n" + "=" * 65)
        print("  有效协议:")
        for name, desc, frame_hex in results:
            print(f"  ✓ {name}: {desc}")
            print(f"    HEX: {frame_hex}")
    else:
        print("\n✗ 没有找到有效协议。")
        print("  可能原因:")
        print("  1. 波特率不是 9600 — 用 --auto-baud 再测")
        print("  2. 发送和接收需不同波特率")
        print('  3. 红外板需要先发送"进入发送模式"的指令')
        print("  4. 换个键值试试 -k 15")


if __name__ == "__main__":
    main()