#!/usr/bin/env python3
# auto_capture_jpeg.py
# 自动连接 ESP32 串口、触发复位、等待多个 JPEG dump、保存并打开所有图片
# 支持新格式：===JPEG_DUMP_BEGIN fmt=<NAME> size=<N>===

import serial
import base64
import re
import sys
import time
import subprocess
from pathlib import Path

PORT = "/dev/cu.usbmodem834401"
BAUD = 115200
OUT_DIR = Path("/Users/rdzleo/Desktop/CogletESP-camera-version/scripts")
TIMEOUT_SEC = 90
MAX_FRAMES = 4

BEGIN_RE = re.compile(r"===JPEG_DUMP_BEGIN\s+(?:fmt=(\S+)\s+)?size=(\d+)===")
END_RE = re.compile(r"===JPEG_DUMP_END===")
B64_RE = re.compile(r"^[A-Za-z0-9+/=]+$")


def main():
    print(f"[·] 打开串口 {PORT} @ {BAUD}")
    ser = serial.Serial(PORT, BAUD, timeout=1)

    print("[·] 复位 ESP32 …")
    ser.dtr = False
    ser.rts = True
    time.sleep(0.1)
    ser.rts = False
    time.sleep(0.1)
    ser.reset_input_buffer()

    print(f"[·] 等待 JPEG_DUMP 标记（最多 {TIMEOUT_SEC}s，期望 {MAX_FRAMES} 张）…")
    start = time.time()
    in_dump = False
    expected_size = 0
    current_fmt = None
    b64_buf = []
    saved_files = []

    while time.time() - start < TIMEOUT_SEC and len(saved_files) < MAX_FRAMES:
        raw = ser.readline()
        if not raw:
            continue
        try:
            line = raw.decode("utf-8", errors="replace").rstrip("\r\n")
        except Exception:
            continue

        if any(k in line for k in ["FaceTracker", "Camera", "panic", "Guru", "ov3660", "Compile time"]):
            print(f"   {line}")

        if not in_dump:
            m = BEGIN_RE.search(line)
            if m:
                in_dump = True
                current_fmt = m.group(1) or "unknown"
                expected_size = int(m.group(2))
                b64_buf = []
                print(f"[+] JPEG_BEGIN fmt={current_fmt} size={expected_size}")
            continue

        if END_RE.search(line):
            in_dump = False
            b64_str = "".join(b64_buf)
            try:
                data = base64.b64decode(b64_str)
            except Exception as e:
                print(f"[!] base64 decode failed: {e}")
                continue

            if len(data) != expected_size:
                print(f"[!] 字节数差异 got={len(data)} expected={expected_size}")

            out_path = OUT_DIR / f"frame_{current_fmt}.jpg"
            out_path.write_bytes(data)
            print(f"[✓] 保存 {out_path.name} ({len(data)/1024:.1f} KB)")
            saved_files.append(out_path)
            continue

        stripped = line.strip()
        if B64_RE.match(stripped):
            b64_buf.append(stripped)

    ser.close()

    if not saved_files:
        print("[!] 没有抓到任何 JPEG 帧")
        return 1

    print(f"\n[✓] 共保存 {len(saved_files)} 张")
    for p in saved_files:
        print(f"    - {p}")
    # 用 Finder 打开目录，用户可以并排对比
    subprocess.run(["open", str(OUT_DIR)])
    # 或者直接打开所有 JPEG
    for p in saved_files:
        subprocess.run(["open", str(p)])
    return 0


if __name__ == "__main__":
    sys.exit(main())