#!/usr/bin/env python3
"""K1 bass/chord register auto-test via J-Link RTT.

Requires firmware with:
  - [PITCH]/[REG] logs (App_Auto.c)
  - RTT cmds: chord key N / chord xpose N / tone pick / log clear

Modes:
  python tools/rtt_pitch_reg_test.py              # inject + live assert
  python tools/rtt_pitch_reg_test.py --listen     # you press keys; we assert
  python tools/rtt_pitch_reg_test.py --analyze oct_reg3.log
"""

from __future__ import annotations

import argparse
import os
import re
import sys
import time
from dataclasses import dataclass, field
from datetime import datetime
from typing import Iterable

DEFAULT_DEVICES = ("Cortex-M4", "AT32F403AC", "AT32F403A")

MIDI_E1 = 28
MIDI_E2 = 40
MIDI_GSHARP4 = 68
XPOSE_FS = 6
BASS_CH_DISP = 8   # 代码/日志通道；BASS_CH=8
DRUM_CH_DISP = 9

PITCH_RE = re.compile(
    r"\[PITCH\]\s+(on|off)\s+ch(\d+)\((\w+)\)\s+"
    r"(\d+)->(\d+)\s+d=(-?\d+)\s+deg=(\d+)\s+chord=(\d+)\s+"
    r"xp=(\d+)\s+xf=(-?\d+)\s+fl=0x([0-9A-Fa-f]+)\s*(.*)$"
)
REG_RE = re.compile(
    r"\[REG\s*\]\s+fp\s+(\d+)->(\d+)\s+chord=(\d+)\s+deg=(\d+)\s+xp=(\d+)\s+xf=(-?\d+)"
)


@dataclass
class PitchEvent:
    is_on: bool
    ch: int
    role: str
    key_in: int
    key_out: int
    delta: int
    deg: int
    chord: int
    xp: int
    xf: int
    flags: int
    tags: str
    raw: str


@dataclass
class CheckResult:
    ok: int = 0
    fail: int = 0
    skip: int = 0
    messages: list[str] = field(default_factory=list)

    def add_ok(self, msg: str) -> None:
        self.ok += 1
        self.messages.append(f"OK  {msg}")

    def add_fail(self, msg: str) -> None:
        self.fail += 1
        self.messages.append(f"FAIL {msg}")

    def add_skip(self, msg: str) -> None:
        self.skip += 1
        self.messages.append(f"SKIP {msg}")


def import_deps() -> None:
    try:
        import pylink  # noqa: F401
    except ImportError as exc:
        raise SystemExit("Missing dependency: pip install pylink-square") from exc


def connect_jlink(device: str):
    import pylink

    jlink = pylink.JLink()
    jlink.open()
    try:
        jlink.exec_command("HideDeviceSelection = 1")
    except Exception:
        pass
    jlink.set_tif(pylink.enums.JLinkInterfaces.SWD)

    last_error = None
    for candidate in (device, *DEFAULT_DEVICES):
        try:
            try:
                jlink.exec_command(f"Device = {candidate}")
            except Exception:
                pass
            jlink.connect(candidate)
            print(f"Connected as {candidate}", flush=True)
            return jlink
        except pylink.errors.JLinkException as exc:
            last_error = exc
    jlink.close()
    raise SystemExit(f"Failed to connect J-Link: {last_error}")


def find_rtt_control_block(jlink, ram_base: int = 0x20000000, ram_size: int = 0x18000) -> int | None:
    needle = b"SEGGER RTT"
    chunk = 0x1000
    was_halted = jlink.halted()
    if not was_halted:
        jlink.halt()
    try:
        for off in range(0, ram_size, chunk):
            data = bytes(jlink.memory_read8(ram_base + off, min(chunk, ram_size - off)))
            idx = data.find(needle)
            if idx >= 0:
                return ram_base + off + idx
    finally:
        if not was_halted:
            if hasattr(jlink, "restart"):
                jlink.restart()
            else:
                jlink.go()
            time.sleep(0.1)
    return None


def wait_rtt_ready(jlink, timeout_s: float = 15.0) -> None:
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        try:
            if jlink.rtt_get_num_up_buffers() > 0:
                return
        except Exception:
            pass
        time.sleep(0.2)
    raise SystemExit("RTT control block found but buffers not ready")


def send_command(jlink, payload: bytes, retries: int = 30) -> None:
    if isinstance(payload, str):
        payload = payload.encode("ascii")
    if not payload.endswith(b"\n"):
        payload += b"\n"
    for _ in range(retries):
        wrote = jlink.rtt_write(0, list(payload))
        if wrote > 0:
            print(f">> {payload.decode('ascii', errors='replace').strip()}", flush=True)
            return
        time.sleep(0.2)
    raise SystemExit(f"Failed to send RTT cmd: {payload!r}")


def drain_rtt(jlink, seconds: float = 0.15) -> str:
    buf = b""
    deadline = time.time() + seconds
    while time.time() < deadline:
        chunk = jlink.rtt_read(0, 4096)
        if chunk:
            buf += bytes(chunk)
        else:
            time.sleep(0.01)
    return buf.decode("utf-8", errors="replace")


def parse_pitch_line(line: str) -> PitchEvent | None:
    m = PITCH_RE.search(line)
    if not m:
        return None
    return PitchEvent(
        is_on=m.group(1) == "on",
        ch=int(m.group(2)),
        role=m.group(3),
        key_in=int(m.group(4)),
        key_out=int(m.group(5)),
        delta=int(m.group(6)),
        deg=int(m.group(7)),
        chord=int(m.group(8)),
        xp=int(m.group(9)),
        xf=int(m.group(10)),
        flags=int(m.group(11), 16),
        tags=m.group(12).strip(),
        raw=line.strip(),
    )


def chord_degree(chord: int) -> int:
    if chord < 1 or chord > 21:
        return 0
    return (chord - 1) // 3 + 1


def need_register_fix(chord: int) -> bool:
    return 7 <= chord <= 21


def map_bass(key: int, xp: int) -> tuple[int, int]:
    """Return (out, flags) mirroring App_Auto_MapBassNote."""
    out = key - 12
    flags = 0x01
    while out < MIDI_E1:
        out += 12
        flags |= 0x02
    if xp >= XPOSE_FS:
        out -= 12
        flags |= 0x04
        while out < MIDI_E1:
            out += 12
            flags |= 0x02
    return max(0, min(127, out)), flags


def map_chord(key: int, xp: int) -> tuple[int, int, bool]:
    """Return (out_or_clamp_hint, flags, need_near).

    When need_near is True, exact out is unknown without chord table;
    caller should only require out in [E2, #G4].
    """
    out = key - 12
    flags = 0x01
    while out < MIDI_E2:
        out += 12
        flags |= 0x08
    need_near = out > MIDI_GSHARP4
    if need_near:
        flags |= 0x10
        out = MIDI_GSHARP4  # placeholder; exact nearest checked soft
    if xp >= XPOSE_FS:
        out -= 12
        flags |= 0x04
        while out < MIDI_E2:
            out += 12
            flags |= 0x08
        if out > MIDI_GSHARP4:
            flags |= 0x10
            need_near = True
            out = MIDI_GSHARP4
    return max(0, min(127, out)), flags, need_near


def expect_for_event(ev: PitchEvent) -> tuple[str, bool]:
    """Return (detail, ok)."""
    if not ev.is_on:
        return "note-off ignored", True

    if ev.delta != ev.key_out - ev.key_in:
        return f"d mismatch {ev.delta} != {ev.key_out - ev.key_in}", False

    deg_exp = chord_degree(ev.chord)
    if deg_exp and ev.deg != deg_exp:
        return f"deg {ev.deg} != expected {deg_exp}", False

    xf_exp = 1 if ev.xp >= XPOSE_FS else 0
    if ev.xf != xf_exp:
        return f"xf {ev.xf} != expected {xf_exp} (xp={ev.xp})", False

    if ev.ch == DRUM_CH_DISP or ev.role == "drum":
        return "drum should not appear in PITCH", False

    fix = need_register_fix(ev.chord)
    if not fix:
        if ev.key_out != ev.key_in or ev.delta != 0:
            return f"I/II should PASS {ev.key_in}->{ev.key_in}, got {ev.key_out}", False
        if "PASS" not in ev.tags and ev.flags != 0:
            # flags may be 0 on pass path
            pass
        return f"PASS I/II {ev.key_in}", True

    if ev.ch == BASS_CH_DISP or ev.role == "bass":
        exp, exp_fl = map_bass(ev.key_in, ev.xp)
        if ev.key_out != exp:
            return f"bass expect {ev.key_in}->{exp}, got {ev.key_out}", False
        if ev.key_out < MIDI_E1:
            return f"bass below E1: {ev.key_out}", False
        # flag bits that must be present
        if (exp_fl & 0x01) and not (ev.flags & 0x01):
            return f"bass missing -12 flag fl=0x{ev.flags:02X}", False
        return f"bass {ev.key_in}->{ev.key_out}", True

    # chord path
    exp, exp_fl, need_near = map_chord(ev.key_in, ev.xp)
    if not (MIDI_E2 <= ev.key_out <= MIDI_GSHARP4):
        return f"chord out {ev.key_out} not in [{MIDI_E2},{MIDI_GSHARP4}]", False
    if need_near:
        if not (ev.flags & 0x10) and "NEAR" not in ev.tags:
            # soft: range ok is enough if nearest not flagged but clamped somehow
            return f"chord NEAR expected (in={ev.key_in} out={ev.key_out})", True
        return f"chord NEAR {ev.key_in}->{ev.key_out}", True
    if ev.key_out != exp:
        return f"chord expect {ev.key_in}->{exp}, got {ev.key_out}", False
    return f"chord {ev.key_in}->{ev.key_out}", True


def analyze_lines(lines: Iterable[str], result: CheckResult | None = None) -> CheckResult:
    result = result or CheckResult()
    pitch_n = 0
    roles: set[str] = set()
    chords: set[int] = set()
    for line in lines:
        line = line.rstrip("\n")
        if "[REG" in line and "fp " in line:
            m = REG_RE.search(line)
            if m:
                chord = int(m.group(3))
                deg = int(m.group(4))
                xp = int(m.group(5))
                xf = int(m.group(6))
                deg_exp = chord_degree(chord)
                xf_exp = 1 if xp >= XPOSE_FS else 0
                if deg_exp and deg != deg_exp:
                    result.add_fail(f"REG deg {deg}!={deg_exp} chord={chord}")
                elif xf != xf_exp:
                    result.add_fail(f"REG xf {xf}!={xf_exp} xp={xp}")
                else:
                    result.add_ok(f"REG fp chord={chord} deg={deg} xp={xp} xf={xf}")

        ev = parse_pitch_line(line)
        if not ev:
            continue
        pitch_n += 1
        roles.add(ev.role)
        chords.add(ev.chord)
        detail, ok = expect_for_event(ev)
        msg = f"ch{ev.ch}({ev.role}) chord={ev.chord} {ev.key_in}->{ev.key_out} | {detail}"
        if ok:
            result.add_ok(msg)
        else:
            result.add_fail(msg)

    if pitch_n == 0:
        result.add_skip("no [PITCH] lines found")
    else:
        if "bass" not in roles and not any(
            f"ch{BASS_CH_DISP}(" in m for m in result.messages
        ):
            result.add_skip(f"no bass PITCH on ch{BASS_CH_DISP} (style may omit bass, or old FW)")
        if not any(need_register_fix(c) for c in chords):
            result.add_skip("no III+ chord keys observed")
        if not any(not need_register_fix(c) and 1 <= c <= 6 for c in chords):
            result.add_skip("no I/II chord keys observed")
    return result


def print_report(result: CheckResult, out_path: str | None = None) -> int:
    print("\n======== PITCH REG TEST REPORT ========", flush=True)
    for msg in result.messages:
        # only print fails + summary skips loudly; OK compacted
        if msg.startswith("FAIL") or msg.startswith("SKIP"):
            print(msg, flush=True)
    ok_short = [m for m in result.messages if m.startswith("OK")]
    print(f"OK={result.ok}  FAIL={result.fail}  SKIP={result.skip}", flush=True)
    if ok_short:
        print(f"(first OK samples: {len(ok_short)} total)", flush=True)
        for m in ok_short[:8]:
            print(m, flush=True)
        if len(ok_short) > 8:
            print(f"... +{len(ok_short) - 8} more OK", flush=True)
    if out_path:
        with open(out_path, "w", encoding="utf-8", newline="\n") as f:
            f.write(f"# pitch reg report {datetime.now().isoformat(timespec='seconds')}\n")
            f.write(f"OK={result.ok} FAIL={result.fail} SKIP={result.skip}\n")
            for m in result.messages:
                f.write(m + "\n")
        print(f"Report saved: {os.path.abspath(out_path)}", flush=True)
    return 0 if result.fail == 0 and result.ok > 0 else 1


class RttSession:
    def __init__(self, jlink):
        self.jlink = jlink
        self.buf = b""
        self.lines: list[str] = []

    def cmd(self, text: str, settle: float = 0.2) -> None:
        send_command(self.jlink, text.encode("ascii"))
        self.pump(settle)

    def pump(self, seconds: float) -> list[str]:
        new_lines: list[str] = []
        deadline = time.time() + seconds
        while time.time() < deadline:
            chunk = self.jlink.rtt_read(0, 4096)
            if chunk:
                self.buf += bytes(chunk)
                while b"\n" in self.buf:
                    raw, self.buf = self.buf.split(b"\n", 1)
                    line = raw.decode("utf-8", errors="replace").rstrip("\r")
                    if line:
                        self.lines.append(line)
                        new_lines.append(line)
                        if (
                            "[PITCH]" in line
                            or "[REG" in line
                            or "[KEY" in line
                            or "CHORD_" in line
                            or "TONE_" in line
                        ):
                            print(line, flush=True)
            else:
                time.sleep(0.01)
        return new_lines

    def cmd_ack(self, text: str, ack: str, timeout_s: float = 3.0, retries: int = 6) -> bool:
        """Send command until ack substring appears (handles RTT down loss under load)."""
        for attempt in range(1, retries + 1):
            # brief quiet drain then send
            self.pump(0.15)
            send_command(self.jlink, text.encode("ascii"))
            deadline = time.time() + timeout_s
            while time.time() < deadline:
                for line in self.pump(0.1):
                    if ack in line:
                        return True
            print(f"!! no ack '{ack}' for '{text}' (try {attempt}/{retries})", flush=True)
            time.sleep(0.25)
        return False

    def stop_band(self) -> None:
        self.cmd_ack("chord key 23", "CHORD_KEY_OK", timeout_s=2.0, retries=4)
        self.pump(0.4)

    def set_xpose(self, xp: int) -> bool:
        return self.cmd_ack(f"chord xpose {xp}", "CHORD_XPOSE_OK", timeout_s=2.0, retries=5)

    def set_key(self, key: int) -> bool:
        return self.cmd_ack(f"chord key {key}", "CHORD_KEY_OK", timeout_s=2.0, retries=5)

    def phase_switch(self, name: str, key: int, xpose: int | None, hold_s: float) -> None:
        """Switch chord/xpose while accompaniment is already running.
        Note: plain `tone pick` clears KEY_ID to 1 — do not call it after inject.
        """
        print(f"\n=== {name}: key={key} xpose={xpose} ===", flush=True)
        if xpose is not None:
            if not self.set_xpose(xpose):
                print(f"FAIL setup xpose={xpose}", flush=True)
        self.set_key(0)
        if not self.set_key(key):
            print(f"FAIL setup key={key}", flush=True)
        self.pump(hold_s)


def run_auto(device: str, hold_s: float, report: str | None, log_path: str | None) -> int:
    import_deps()
    jlink = connect_jlink(device)
    try:
        cb = find_rtt_control_block(jlink)
        if cb is None:
            raise SystemExit("SEGGER 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)

        sess.cmd("log clear", 0.3)
        sess.cmd("tone local", 1.0)
        # Start band first (pick forces chord=1), then inject keys while playing
        if not sess.set_xpose(2):
            print("FAIL initial xpose=2", flush=True)
        sess.cmd("tone pick", settle=0.8)
        print("\n=== I (default after pick) ===", flush=True)
        sess.pump(hold_s)

        sess.phase_switch("I/II key2 PASS", key=2, xpose=None, hold_s=hold_s)
        sess.phase_switch("III key8 -12", key=8, xpose=None, hold_s=hold_s)
        sess.phase_switch("VII key20 -12", key=20, xpose=None, hold_s=hold_s)
        sess.phase_switch("VII key20 xpose=#F", key=20, xpose=6, hold_s=hold_s)

        print("\n=== stop ===", flush=True)
        sess.stop_band()
        sess.pump(0.4)

        if log_path:
            with open(log_path, "w", encoding="utf-8", newline="\n") as f:
                f.write(f"# auto pitch test {datetime.now().isoformat(timespec='seconds')}\n")
                for line in sess.lines:
                    f.write(line + "\n")
            print(f"Log saved: {os.path.abspath(log_path)}", flush=True)

        joined = "\n".join(sess.lines)
        if "CHORD_KEY_OK" not in joined and "CHORD_KEY_BAD" not in joined:
            print(
                "WARNING: no CHORD_KEY_OK — firmware may lack 'chord key' RTT cmd. "
                "Rebuild/flash, or use --listen / --analyze.",
                flush=True,
            )

        result = analyze_lines(sess.lines)
        has_pass = any("PASS I/II" in m for m in result.messages)
        has_iii_raw = any(
            "[PITCH]" in line and ("chord=8" in line or "chord=20" in line or "deg=3" in line or "deg=7" in line)
            for line in sess.lines
        )
        has_iii_ok = any(
            m.startswith("OK") and "PASS" not in m and "(chord)" in m
            for m in result.messages
        )
        has_xf = any("[PITCH]" in line and "xf=1" in line and "deg=7" in line for line in sess.lines)
        if not has_pass:
            result.add_fail("coverage: no I/II PASS samples")
        if not has_iii_raw:
            result.add_fail("coverage: no III+/VII chord in PITCH")
        elif not has_iii_ok:
            result.add_fail("coverage: III+ present but mapping asserts failed")
        if not has_xf:
            result.add_skip("coverage: no VII+xf=1 samples")

        return print_report(result, report)
    finally:
        try:
            jlink.rtt_stop()
        except Exception:
            pass
        jlink.close()


def run_listen(device: str, seconds: float, report: str | None, log_path: str | None) -> int:
    import_deps()
    jlink = connect_jlink(device)
    try:
        cb = find_rtt_control_block(jlink)
        if cb is None:
            raise SystemExit("SEGGER 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.2)
        sess.cmd("log clear", 0.2)
        print(
            f"Listening {seconds:.0f}s — press I/II then III+ chords; set transpose>=#F if possible.",
            flush=True,
        )
        sess.pump(seconds)
        if log_path:
            with open(log_path, "w", encoding="utf-8", newline="\n") as f:
                f.write(f"# listen pitch test {datetime.now().isoformat(timespec='seconds')}\n")
                for line in sess.lines:
                    f.write(line + "\n")
            print(f"Log saved: {os.path.abspath(log_path)}", flush=True)
        result = analyze_lines(sess.lines)
        return print_report(result, report)
    finally:
        try:
            jlink.rtt_stop()
        except Exception:
            pass
        jlink.close()


def run_analyze(path: str, report: str | None) -> int:
    with open(path, encoding="utf-8", errors="replace") as f:
        lines = f.readlines()
    result = analyze_lines(lines)
    return print_report(result, report)


def main() -> None:
    parser = argparse.ArgumentParser(description="K1 pitch/register RTT auto-test")
    parser.add_argument("--device", default="AT32F403AC")
    parser.add_argument("--hold", type=float, default=2.5, help="Seconds to hold each injected key")
    parser.add_argument("--seconds", type=float, default=45.0, help="Listen duration")
    parser.add_argument("--listen", action="store_true", help="Do not inject; monitor only")
    parser.add_argument("--analyze", metavar="LOG", help="Offline analyze a dump log")
    parser.add_argument("--out-log", default="", help="Save captured RTT lines")
    parser.add_argument("--report", default="", help="Save pass/fail report")
    args = parser.parse_args()

    report = args.report or None
    log_path = args.out_log or None

    if args.analyze:
        raise SystemExit(run_analyze(args.analyze, report))
    if args.listen:
        raise SystemExit(run_listen(args.device, args.seconds, report, log_path))
    raise SystemExit(run_auto(args.device, args.hold, report, log_path))


if __name__ == "__main__":
    main()