#!/usr/bin/env python3
"""Probe 1.bin (rhythm) and 3.bin (universal) for bass ch8 + plan compliance.

Requires FW with: tone bin1 / tone bin3 / tone start / chord key / chord xpose / PITCH ch8(bass)
"""

from __future__ import annotations

import argparse
import os
import re
import sys
import time
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rtt_pitch_reg_test import (  # noqa: E402
    RttSession,
    analyze_lines,
    connect_jlink,
    expect_for_event,
    find_rtt_control_block,
    import_deps,
    parse_pitch_line,
    wait_rtt_ready,
)

PITCH_RE_CH = re.compile(r"\[PITCH\].*?ch(\d+)\((\w+)\)")
CH_SEEN_RE = re.compile(r"ch-seen ch(\d+) role=(\w+) key=(-?\d+)")
TONE_OK_RE = re.compile(r"TONE_BIN([13]).*name=(\S+).*count=(\d+)")


@dataclass
class PresetResult:
    bin_id: str
    idx: int
    name: str = ""
    channels: Counter = field(default_factory=Counter)
    roles: Counter = field(default_factory=Counter)
    bass_n: int = 0
    drum_n: int = 0
    chord_n: int = 0
    pass_ok: int = 0
    pass_fail: int = 0
    iii_ok: int = 0
    iii_fail: int = 0
    xf_ok: int = 0
    xf_fail: int = 0
    bass_ok: int = 0
    bass_fail: int = 0
    samples: list[str] = field(default_factory=list)


def run_preset(sess: RttSession, bin_id: str, idx: int, hold: float) -> PresetResult:
    res = PresetResult(bin_id=bin_id, idx=idx)
    mark = len(sess.lines)

    if bin_id == "1":
        ok = sess.cmd_ack(f"tone bin1 {idx}", "TONE_BIN1", timeout_s=3.0, retries=4)
    else:
        ok = sess.cmd_ack(f"tone bin3 {idx}", "TONE_BIN3", timeout_s=3.0, retries=4)
    if not ok:
        res.samples.append("FAIL load ack")
        return res

    # parse name from recent lines
    for line in sess.lines[-8:]:
        m = TONE_OK_RE.search(line)
        if m and m.group(1) == bin_id:
            res.name = m.group(2)
            break

    sess.set_xpose(2)
    sess.cmd("tone start", settle=0.7)

    # I/II
    sess.set_key(0)
    sess.set_key(2)
    sess.pump(hold)

    # III
    sess.set_key(0)
    sess.set_key(8)
    sess.pump(hold)

    # VII + #F
    sess.set_xpose(6)
    sess.set_key(0)
    sess.set_key(20)
    sess.pump(hold)

    sess.set_key(23)  # stop
    sess.pump(0.3)

    chunk = sess.lines[mark:]
    for line in chunk:
        m = CH_SEEN_RE.search(line)
        if m:
            ch, role = int(m.group(1)), m.group(2)
            res.channels[ch] += 1
            res.roles[role] += 1
            res.samples.append(line.strip())
        ev = parse_pitch_line(line)
        if not ev or not ev.is_on:
            continue
        res.channels[ev.ch] += 1
        res.roles[ev.role] += 1
        if ev.role == "bass" or ev.ch == 8:
            res.bass_n += 1
        elif ev.role == "drum" or ev.ch == 9:
            res.drum_n += 1
        else:
            res.chord_n += 1

        detail, ok = expect_for_event(ev)
        tag = f"ch{ev.ch}({ev.role}) {ev.key_in}->{ev.key_out} chord={ev.chord} | {detail}"
        if ev.chord <= 6:
            if ok:
                res.pass_ok += 1
            else:
                res.pass_fail += 1
                res.samples.append("FAIL " + tag)
        elif ev.xf:
            if ok:
                res.xf_ok += 1
            else:
                res.xf_fail += 1
                res.samples.append("FAIL " + tag)
        elif ev.role == "bass" or ev.ch == 8:
            if ok:
                res.bass_ok += 1
            else:
                res.bass_fail += 1
                res.samples.append("FAIL " + tag)
        else:
            if ok:
                res.iii_ok += 1
            else:
                res.iii_fail += 1
                res.samples.append("FAIL " + tag)

    return res


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

    ap = argparse.ArgumentParser()
    ap.add_argument("--hold", type=float, default=2.2)
    ap.add_argument("--bin1", default="0,1,2,3,4,5,8,12,16,20,24,28", help="1.bin indices")
    ap.add_argument("--bin3", default="0,1,2", help="3.bin indices")
    ap.add_argument("--out", default="")
    args = ap.parse_args()

    bin1_idxs = [int(x) for x in args.bin1.split(",") if x.strip() != ""]
    bin3_idxs = [int(x) for x in args.bin3.split(",") if x.strip() != ""]

    import_deps()
    jlink = connect_jlink("AT32F403AC")
    results: list[PresetResult] = []
    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)
        sess.cmd("log clear", 0.3)

        print("\n######## BIN1 / 1.bin rhythms ########", flush=True)
        for idx in bin1_idxs:
            print(f"\n--- BIN1 idx={idx} ---", flush=True)
            r = run_preset(sess, "1", idx, args.hold)
            results.append(r)
            print(
                f"name={r.name} ch={dict(r.channels)} roles={dict(r.roles)} "
                f"bass={r.bass_n} drum={r.drum_n} chord={r.chord_n} "
                f"I/II {r.pass_ok}/{r.pass_fail} III {r.iii_ok}/{r.iii_fail} "
                f"xf {r.xf_ok}/{r.xf_fail} bassMap {r.bass_ok}/{r.bass_fail}",
                flush=True,
            )

        print("\n######## BIN3 / 3.bin universal ########", flush=True)
        for idx in bin3_idxs:
            print(f"\n--- BIN3 idx={idx} ---", flush=True)
            r = run_preset(sess, "3", idx, args.hold)
            results.append(r)
            print(
                f"name={r.name} ch={dict(r.channels)} roles={dict(r.roles)} "
                f"bass={r.bass_n} drum={r.drum_n} chord={r.chord_n} "
                f"I/II {r.pass_ok}/{r.pass_fail} III {r.iii_ok}/{r.iii_fail} "
                f"xf {r.xf_ok}/{r.xf_fail} bassMap {r.bass_ok}/{r.bass_fail}",
                flush=True,
            )

        # Summary vs plan
        print("\n======== PLAN CHECK (bass和弦分通道八度) ========", flush=True)
        with_bass = [r for r in results if r.bass_n > 0]
        print(f"Presets tested: {len(results)}; with ch8(bass) NoteOn: {len(with_bass)}", flush=True)
        if with_bass:
            for r in with_bass:
                print(
                    f"  BASS HIT {r.bin_id}.bin[{r.idx}] {r.name}: "
                    f"bass_n={r.bass_n} map_ok={r.bass_ok} fail={r.bass_fail}",
                    flush=True,
                )
        else:
            print("  NO preset produced ch8(bass). Plan bass rule NOT verified on-device.", flush=True)

        # Aggregate chord rules
        pass_ok = sum(r.pass_ok for r in results)
        pass_fail = sum(r.pass_fail for r in results)
        iii_ok = sum(r.iii_ok for r in results)
        iii_fail = sum(r.iii_fail for r in results)
        xf_ok = sum(r.xf_ok for r in results)
        xf_fail = sum(r.xf_fail for r in results)
        bass_ok = sum(r.bass_ok for r in results)
        bass_fail = sum(r.bass_fail for r in results)

        def gate(name, ok, fail, need=True):
            status = "PASS" if fail == 0 and (ok > 0 or not need) else ("FAIL" if fail else "SKIP")
            print(f"  [{status}] {name}: ok={ok} fail={fail}", flush=True)
            return status != "FAIL"

        all_ok = True
        all_ok &= gate("I/II PASS (keys1-6)", pass_ok, pass_fail)
        all_ok &= gate("III+ chord map", iii_ok, iii_fail)
        all_ok &= gate("xpose>=#F extra-12", xf_ok, xf_fail)
        all_ok &= gate("bass ch8 map", bass_ok, bass_fail, need=False)
        if not with_bass:
            all_ok = False
            print("  [FAIL] plan requires bass on code ch8 — never seen across 1.bin/3.bin samples", flush=True)

        # channel histogram
        ch_all: Counter = Counter()
        for r in results:
            ch_all.update(r.channels)
        print(f"Channel histogram: {dict(sorted(ch_all.items()))}", flush=True)

        stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        out = args.out or os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            f"bin13_reg_{stamp}.log",
        )
        with open(out, "w", encoding="utf-8", newline="\n") as f:
            f.write(f"# bin1/bin3 plan check {datetime.now().isoformat(timespec='seconds')}\n")
            for line in sess.lines:
                f.write(line + "\n")
            f.write("\n# SUMMARY\n")
            for r in results:
                f.write(
                    f"# {r.bin_id}[{r.idx}] {r.name} ch={dict(r.channels)} "
                    f"bass={r.bass_n} I={r.pass_ok}/{r.pass_fail} "
                    f"III={r.iii_ok}/{r.iii_fail} xf={r.xf_ok}/{r.xf_fail} "
                    f"bassMap={r.bass_ok}/{r.bass_fail}\n"
                )
        print(f"Log: {out}", flush=True)
        return 0 if all_ok else 1
    finally:
        try:
            jlink.rtt_stop()
        except Exception:
            pass
        jlink.close()


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