#!/usr/bin/env python3
# extract_jpeg_from_log.py
# 从 ESP32 串口日志中提取 base64 编码的 JPEG 图像并保存为 .jpg 文件
#
# 用法：
#   1) 启动 monitor 并把日志重定向到文件：
#        idf.py -p /dev/cu.usbmodem834401 monitor > /tmp/esp32.log
#      或者直接从已有日志提取：
#        python3 extract_jpeg_from_log.py /tmp/esp32.log
#   2) 设备启动后会打印 "===JPEG_DUMP_BEGIN===" .... "===JPEG_DUMP_END==="
#   3) 运行此脚本会在当前目录生成 frame_001.jpg 等文件
#
# 也支持从 stdin 实时读取：
#        idf.py monitor | python3 extract_jpeg_from_log.py -

import sys
import re
import base64
from pathlib import Path

BEGIN_RE = re.compile(r"===JPEG_DUMP_BEGIN\s+size=(\d+)===")
END_RE = re.compile(r"===JPEG_DUMP_END===")
# 匹配纯 base64 行（不含普通文本）
B64_RE = re.compile(r"^[A-Za-z0-9+/=]+$")


def extract(lines, out_dir=Path(".")):
    frame_idx = 0
    in_dump = False
    b64_buf = []
    expected_size = 0

    for raw in lines:
        line = raw.rstrip("\r\n")
        # 去掉 ESP 日志时间戳前缀可能引入的干扰：只在 begin/end 标记附近处理
        if not in_dump:
            m = BEGIN_RE.search(line)
            if m:
                in_dump = True
                expected_size = int(m.group(1))
                b64_buf = []
                print(f"[+] JPEG_BEGIN 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"[!] size mismatch: got {len(data)} expected {expected_size} "
                    f"(可能是 monitor 丢字节，仍尝试保存)"
                )
            frame_idx += 1
            out_path = out_dir / f"frame_{frame_idx:03d}.jpg"
            out_path.write_bytes(data)
            print(f"[✓] saved {out_path} ({len(data)} bytes)")
            # macOS: 自动打开
            if sys.platform == "darwin":
                import subprocess
                subprocess.run(["open", str(out_path)])
            continue

        # 处于 dump 区间，只收集看起来是 base64 的行
        stripped = line.strip()
        if B64_RE.match(stripped):
            b64_buf.append(stripped)
        # 其他行（比如 "I (xxx) TAG:" 的日志）忽略


def main():
    if len(sys.argv) < 2:
        print("Usage: extract_jpeg_from_log.py <logfile | ->")
        sys.exit(1)
    src = sys.argv[1]
    if src == "-":
        extract(sys.stdin)
    else:
        with open(src, "r", errors="replace") as f:
            extract(f)


if __name__ == "__main__":
    main()