#!/usr/bin/env python3
"""Стенд декодирования FT4: аудио TCI против виртуального кабеля, бок о бок.

Зачем: MSHV по виртуальному аудиокабелю FT4 декодирует, а по TCI — нет, хотя
источник звука один и тот же. Стенд снимает оба потока ОДНОВРЕМЕННО, режет их
на нарезки FT4 по единым часам и отдаёт каждую нарезку настоящему декодеру
(`jt9 --ft4` из WSJT-X). Дальше остаётся смотреть в таблицу, а не гадать.

Три колонки не случайны:

  cable  — эталон: то, что слышит MSHV сейчас (PipeWire-источник, обычно
           `wsjtx_tx_source`);
  tci    — наш поток как есть, включая обрезание по ±1.0 при укладке в int16
           (ровно то, что делает клиент, попросивший int16);
  tci-n  — тот же поток, но нормированный по пику перед укладкой в WAV.

Если декодируется `tci-n`, а `tci` — нет, значит дело в УРОВНЕ: тап RX_AUDIO
снят до регулятора громкости, и пики WDSP выходят за ±1.0 (в этом проекте они
доходили до 2.07). Всё, что выше единицы, у клиента срезается — сигнал в
водопаде становится шире и грязнее, а декодер разваливается. Если не
декодируется ни то, ни другое, а cable декодируется — смотреть надо на
непрерывность потока (потерянные блоки) и на частоту дискретизации.

Стенд ТОЛЬКО СЛУШАЕТ: ни одной команды, способной тронуть передачу.

Запуск (приложение работает, TCI поднят, в MSHV идёт приём FT4):
    test/tci/ft4_bench.py --rx 1 --windows 4
    test/tci/ft4_bench.py --rx 0 --device wsjtx_tx_source --keep
"""

import argparse
import math
import os
import struct
import subprocess
import sys
import tempfile
import threading
import time
import wave

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from live_rx_test import WS   # noqa: E402

HDR = struct.Struct("<16I")
STREAM_RX_AUDIO = 1
FMT_FLOAT32 = 3
RATE = 12000          # то, что ждёт jt9
FT4_PERIOD = 7.5      # длина нарезки FT4, с
FT4_SAMPLES = int(RATE * FT4_PERIOD)


# ── съём звука ───────────────────────────────────────────────────────────
class TciCapture(threading.Thread):
    """RX-аудио одного приёмника TCI: float32, моно, 12 кГц.

    Просим именно float32 — иначе истинный уровень не увидеть: в int16 сервер
    сам зажимает всё, что выше ±1.0, и стенд мерил бы уже обрезанное.
    """

    def __init__(self, host, port, rx, channels=2, like_mshv=True):
        super().__init__(daemon=True)
        self.ws = WS(host, port)
        self.rx = rx
        self.channels = channels
        # ★Разбирать блок КАК MSHV: он берёт число байт из поля length
        # (`cr2 = pStream->length*bit_s`) и шагает кадрами `chan*bit_s`, то
        # есть считает length вещественными отсчётами ВСЕГО блока. Если сервер
        # понимает length как «на канал», у стерео разбирается ровно половина
        # блока — и звук у клиента идёт с дырами. Режим --raw читает всю
        # полезную нагрузку, игнорируя length: так видно, что данные-то целы.
        self.like_mshv = like_mshv
        self.short_blocks = 0
        self.samples = []      # float
        self.t0 = None         # время первого отсчёта, monotonic
        self.blocks = 0
        self.bad_rate = 0
        self.stop_flag = False
        self.lock = threading.Lock()

    def start_stream(self):
        self.ws.send_text(f"audio_samplerate:{RATE};")
        self.ws.send_text("audio_stream_sample_type:float32;")
        self.ws.send_text(f"audio_stream_channels:{self.channels};")
        self.ws.send_text("audio_stream_samples:512;")
        self.ws.send_text(f"audio_start:{self.rx};")

    def run(self):
        while not self.stop_flag:
            frame = self.ws.recv_frame(0.3)
            if not frame:
                continue
            op, payload = frame
            if op != 2 or len(payload) < HDR.size:
                continue
            head = HDR.unpack(payload[:HDR.size])
            receiver, srate, fmt, length = head[0], head[1], head[2], head[5]
            stype, chans = head[6], head[7]
            if stype != STREAM_RX_AUDIO or receiver != self.rx:
                continue
            if srate != RATE:
                self.bad_rate += 1
                continue
            data = payload[HDR.size:]
            if self.like_mshv:
                want = int(length) * 4           # cr2 = length * bit_s
                if want < len(data):
                    self.short_blocks += 1
                data = data[:min(want, len(data))]
            n = len(data) // 4
            vals = struct.unpack(f"<{n}f", data[:n * 4])
            if chans == 2:                       # берём левый канал
                vals = vals[0::2]
            now = time.monotonic()
            with self.lock:
                if self.t0 is None:
                    self.t0 = now - len(vals) / RATE
                self.samples.extend(vals)
                self.blocks += 1
                _ = length

    def close(self):
        self.stop_flag = True
        try:
            self.ws.send_text(f"audio_stop:{self.rx};")
            time.sleep(0.2)
        except OSError:
            pass
        self.ws.close()


class CableCapture(threading.Thread):
    """Виртуальный кабель через parec: моно, 12 кГц, s16 — эталонная дорожка."""

    def __init__(self, device):
        super().__init__(daemon=True)
        self.proc = subprocess.Popen(
            ["parec", f"--device={device}", "--format=s16le",
             f"--rate={RATE}", "--channels=1", "--latency-msec=50"],
            stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
        self.samples = []
        self.t0 = None
        self.stop_flag = False
        self.lock = threading.Lock()

    def run(self):
        while not self.stop_flag:
            block = self.proc.stdout.read(4096)
            if not block:
                break
            n = len(block) // 2
            vals = struct.unpack(f"<{n}h", block[:n * 2])
            now = time.monotonic()
            with self.lock:
                if self.t0 is None:
                    self.t0 = now - n / RATE
                self.samples.extend(v / 32768.0 for v in vals)

    def close(self):
        self.stop_flag = True
        try:
            self.proc.terminate()
            self.proc.wait(timeout=2)
        except Exception:
            self.proc.kill()


# ── нарезка, WAV, декодер ────────────────────────────────────────────────
def cut(cap, t_start, count):
    """Отсчёты, начинающиеся в момент t_start (monotonic). None — не хватило."""
    with cap.lock:
        if cap.t0 is None:
            return None
        first = int(round((t_start - cap.t0) * RATE))
        if first < 0 or first + count > len(cap.samples):
            return None
        return list(cap.samples[first:first + count])


def stats(vals):
    if not vals:
        return 0.0, 0.0, 0.0
    peak = max(abs(v) for v in vals)
    rms = math.sqrt(sum(v * v for v in vals) / len(vals))
    over = sum(1 for v in vals if abs(v) >= 1.0)
    return peak, rms, 100.0 * over / len(vals)


def write_wav(path, vals, normalize):
    scale = 1.0
    if normalize:
        peak = max((abs(v) for v in vals), default=0.0)
        if peak > 1e-9:
            scale = 0.9 / peak
    frames = bytearray()
    for v in vals:
        s = int(round(v * scale * 32767.0))
        frames += struct.pack("<h", max(-32768, min(32767, s)))
    with wave.open(path, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(RATE)
        w.writeframes(bytes(frames))


def decode(path, workdir):
    """jt9 --ft4. Возвращает список строк декодера."""
    try:
        res = subprocess.run(
            ["jt9", "--ft4", "-p", "7.5", "-d", "3", "-L", "100", "-H", "3000",
             "-a", workdir, "-t", workdir, os.path.abspath(path)],
            cwd=workdir, capture_output=True, text=True, timeout=60)
    except FileNotFoundError:
        print("jt9 не найден — поставьте WSJT-X", file=sys.stderr)
        sys.exit(2)
    except subprocess.TimeoutExpired:
        return ["<таймаут декодера>"]
    out = []
    for line in res.stdout.splitlines():
        line = line.strip()
        if not line or line.startswith("<Decode"):
            continue
        out.append(line)
    return out


def main():
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=40001)
    ap.add_argument("--rx", type=int, default=1,
                    help="приёмник TCI (1 = «TCI Client rx2» в MSHV)")
    ap.add_argument("--device", default="wsjtx_tx_source",
                    help="источник PipeWire с тем же звуком (виртуальный кабель)")
    ap.add_argument("--windows", type=int, default=4,
                    help="сколько нарезок FT4 снять (по 7.5 с)")
    ap.add_argument("--channels", type=int, default=2, choices=(1, 2),
                    help="каналов в потоке TCI (2 = как просит MSHV)")
    ap.add_argument("--raw", action="store_true",
                    help="читать всю полезную нагрузку, игнорируя поле length "
                         "(MSHV так НЕ умеет — см. шапку)")
    ap.add_argument("--no-cable", action="store_true",
                    help="без эталона: только поток TCI")
    ap.add_argument("--keep", action="store_true",
                    help="оставить WAV-файлы и назвать каталог")
    args = ap.parse_args()

    workdir = tempfile.mkdtemp(prefix="ft4bench-")
    print(f"Стенд FT4: приёмник TCI {args.rx}"
          + ("" if args.no_cable else f", эталон — {args.device}"))
    print(f"рабочий каталог: {workdir}")

    tci = TciCapture(args.host, args.port, args.rx, args.channels,
                     like_mshv=not args.raw)
    cable = None
    try:
        tci.start()
        time.sleep(0.5)
        tci.start_stream()
        if not args.no_cable:
            cable = CableCapture(args.device)
            cable.start()

        # Ждём, пока пойдут оба потока: иначе первая нарезка окажется пустой.
        deadline = time.monotonic() + 5
        while time.monotonic() < deadline:
            if tci.t0 is not None and (cable is None or cable.t0 is not None):
                break
            time.sleep(0.1)
        if tci.t0 is None:
            print("  FAIL блоков RX-аудио от TCI нет вовсе"
                  + (f" (заголовков с чужой частотой: {tci.bad_rate})"
                     if tci.bad_rate else ""))
            return 1
        if cable is not None and cable.t0 is None:
            print(f"  FAIL с источника {args.device} звук не идёт")
            return 1

        print(f"{'нарезка':>8} | {'cable':>5} | {'tci':>5} | {'tci-n':>5} | "
              f"уровень TCI (пик / RMS / за ±1.0)")
        print("-" * 78)

        # Нарезки FT4 идут по единым часам: старт кратен 7.5 с от начала минуты.
        wall = time.time()
        mono = time.monotonic()
        n_ok = {"cable": 0, "tci": 0, "tcin": 0}
        done = 0
        # ★Потолок по времени: если поток идёт медленнее реального времени
        # (клиент разбирает не весь блок), нарезка не наберётся НИКОГДА, и без
        # этого потолка стенд просто висел бы. Само по себе срабатывание —
        # диагноз, поэтому ниже печатается темп набора.
        started = time.monotonic()
        limit = started + args.windows * FT4_PERIOD * 3 + 25
        while done < args.windows:
            if time.monotonic() > limit:
                with tci.lock:
                    got = len(tci.samples) / RATE
                elapsed = time.monotonic() - started
                print("-" * 78)
                print(f"  FAIL нарезки не набираются: за {elapsed:.1f} с "
                      f"реального времени поток TCI дал {got:.1f} с звука "
                      f"({100.0 * got / max(elapsed, 0.1):.0f}% темпа)")
                if tci.short_blocks:
                    print(f"  ★и {tci.short_blocks} блоков разобраны НЕ ЦЕЛИКОМ: "
                          f"поле length меньше присланных данных — клиент вроде "
                          f"MSHV теряет остаток каждого блока")
                break
            now_wall = time.time()
            seq = math.floor(now_wall / FT4_PERIOD) * FT4_PERIOD
            t_start_wall = seq - FT4_PERIOD          # предыдущая, уже целая
            t_start = mono + (t_start_wall - wall)
            # Ждём, пока накопится хвост нарезки у обоих потоков.
            block_tci = cut(tci, t_start, FT4_SAMPLES)
            block_cab = None if cable is None else cut(cable, t_start, FT4_SAMPLES)
            if block_tci is None or (cable is not None and block_cab is None):
                time.sleep(0.3)
                continue

            label = time.strftime("%H:%M:%S", time.localtime(t_start_wall))
            peak, rms, over = stats(block_tci)

            f_tci = os.path.join(workdir, f"tci_{label.replace(':', '')}.wav")
            f_tcn = os.path.join(workdir, f"tcin_{label.replace(':', '')}.wav")
            write_wav(f_tci, block_tci, normalize=False)
            write_wav(f_tcn, block_tci, normalize=True)
            d_tci = decode(f_tci, workdir)
            d_tcn = decode(f_tcn, workdir)
            d_cab = []
            if block_cab is not None:
                f_cab = os.path.join(workdir, f"cable_{label.replace(':', '')}.wav")
                write_wav(f_cab, block_cab, normalize=False)
                d_cab = decode(f_cab, workdir)

            n_ok["cable"] += len(d_cab)
            n_ok["tci"] += len(d_tci)
            n_ok["tcin"] += len(d_tcn)
            print(f"{label:>8} | {len(d_cab):>5} | {len(d_tci):>5} | "
                  f"{len(d_tcn):>5} | пик {peak:.3f}  RMS {rms:.3f}  "
                  f"{over:.2f}%")
            for src, lines in (("cable", d_cab), ("tci", d_tci), ("tci-n", d_tcn)):
                for line in lines[:3]:
                    print(f"           {src:>5}: {line}")
            done += 1
            # Следующая нарезка появится не раньше чем через период.
            time.sleep(FT4_PERIOD * 0.9)

        print("-" * 78)
        print(f"Итого декодов: cable {n_ok['cable']}, tci {n_ok['tci']}, "
              f"tci-n {n_ok['tcin']}")
        if tci.short_blocks:
            print(f"★блоков, у которых length меньше присланных данных: "
                  f"{tci.short_blocks} — клиент вроде MSHV разбирает только "
                  f"часть блока, остальное теряется")
        if n_ok["tci"] == 0 and n_ok["tcin"] > 0:
            print("⇒ дело в УРОВНЕ: поток TCI режется по ±1.0 (см. шапку стенда)")
        elif n_ok["tci"] == 0 and n_ok["cable"] > 0:
            print("⇒ уровень ни при чём: смотреть непрерывность потока и частоту")
    finally:
        tci.close()
        if cable is not None:
            cable.close()
        if not args.keep:
            for name in os.listdir(workdir):
                try:
                    os.unlink(os.path.join(workdir, name))
                except OSError:
                    pass
            try:
                os.rmdir(workdir)
            except OSError:
                pass
    return 0


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