#!/usr/bin/env python3
"""Scan K1 I2C1 via RTT and report amp (0x30) / mic (0x5A)."""

from __future__ import annotations

import os
import re
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from rtt_pitch_reg_test import (  # noqa: E402
    RttSession,
    connect_jlink,
    find_rtt_control_block,
    import_deps,
    wait_rtt_ready,
)

ACK_RE = re.compile(r"^I2C ACK 0x([0-9A-Fa-f]{2})(?:\s+(\S+))?")
END_RE = re.compile(r"^I2C_SCAN_END n=(\d+)")
AMP_ADDR = 0x30
MIC_ADDR = 0x5A


def main() -> int:
    if hasattr(sys.stdout, "reconfigure"):
        try:
            sys.stdout.reconfigure(encoding="utf-8", errors="replace")
        except Exception:
            pass

    import_deps()
    jlink = connect_jlink("AT32F403AC")
    try:
        cb = find_rtt_control_block(jlink)
        if cb is None:
            raise SystemExit("RTT CB not found")
        print(f"RTT CB @ 0x{cb:08X}", flush=True)
        jlink.rtt_start(cb)
        wait_rtt_ready(jlink)

        sess = RttSession(jlink)
        sess.pump(0.3)

        found: dict[int, str] = {}
        end_n: int | None = None

        def consume(lines: list[str]) -> None:
            nonlocal end_n
            for line in lines:
                text = line.strip()
                if not text:
                    continue
                if text.startswith("I2C "):
                    print(text, flush=True)
                m = ACK_RE.match(text)
                if m:
                    found[int(m.group(1), 16)] = m.group(2) or ""
                    continue
                m = END_RE.match(text)
                if m:
                    end_n = int(m.group(1))

        for attempt in range(1, 5):
            mark = len(sess.lines)
            sess.cmd("i2c scan", settle=0.4)
            deadline = time.time() + 4.0
            while time.time() < deadline:
                consume(sess.lines[mark:])
                mark = len(sess.lines)
                if end_n is not None:
                    break
                sess.pump(0.15)
            consume(sess.lines[mark:])
            if end_n is not None:
                break
            print(f"!! no I2C_SCAN_END (try {attempt}/4)", flush=True)

        if end_n is None:
            print("FAIL: firmware did not answer I2C_SCAN_END", flush=True)
            return 1

        others = [a for a in sorted(found) if a not in (AMP_ADDR, MIC_ADDR)]
        amp_ok = AMP_ADDR in found
        mic_ok = MIC_ADDR in found

        print("", flush=True)
        print("======== I2C SCAN ========", flush=True)
        print(
            f"功放 AD85050  7bit 0x{AMP_ADDR:02X} / 写 0x{AMP_ADDR << 1:02X}: "
            f"{'FOUND' if amp_ok else 'MISSING'}",
            flush=True,
        )
        print(
            f"麦克风 AW816  7bit 0x{MIC_ADDR:02X} / 写 0x{MIC_ADDR << 1:02X}: "
            f"{'FOUND' if mic_ok else 'MISSING'}",
            flush=True,
        )
        if others:
            extra = " ".join(f"0x{a:02X}" for a in others)
            print(f"其它 ACK: {extra}", flush=True)
        else:
            print("其它 ACK: (无)", flush=True)
        print(f"I2C_SCAN_END n={end_n}  unique={len(found)}", flush=True)
        return 0 if amp_ok and mic_ok else 2
    finally:
        try:
            jlink.rtt_stop()
        except Exception:
            pass
        jlink.close()


if __name__ == "__main__":
    raise SystemExit(main())