"""Plays pre-recorded WAV clips via AudioClient.PlayStream — bypasses TtsMaker.

Expected WAV format (required by the G1 audio channel):
    16 kHz, mono, 16-bit signed PCM.

Library layout under assets/audio/:
    fixed/<key>.wav              e.g. safe.wav, unsafe_generic.wav, ready.wav
    unsafe_missing/<key>.wav     e.g. helmet.wav, vest.wav, helmet_vest.wav
                                  (key is sorted-joined PPE names, "_" separator)

Callers look up by (category, key). Missing clip → returns False so the caller
can fall back to TtsMaker.
"""
from __future__ import annotations

import datetime
import json
import threading
import time
import wave
from pathlib import Path
from typing import Dict, Optional, Tuple

from core.paths import PROJECT_ROOT

AUDIO_ROOT = PROJECT_ROOT / "assets" / "audio"


def _ts() -> str:
    return datetime.datetime.now().strftime("%H:%M:%S.%f")[:-3]

EXPECTED_RATE     = 16000
EXPECTED_CHANNELS = 1
EXPECTED_WIDTH    = 2          # bytes per sample (int16)
PLAY_CHUNK_BYTES  = 96000      # 3 s per PlayStream call (matches the Unitree example)
PLAY_APP_NAME     = "saqr_audio"

# The G1 arm action and audio stack share a firmware busy state. If we try to
# PlayStream while an arm action is still being processed, chunk 0 is often
# rejected with rc=3104. Later chunks can fail too if the firmware hasn't
# fully processed the previous chunk. Retry both cases with back-off.
CHUNK0_RETRIES     = 4
CHUNK0_BACKOFF_S   = 1.0    # 1,2,3,4 s — total ≈ 10 s, covers a full arm cycle
CHUNKN_RETRIES     = 2
CHUNKN_BACKOFF_S   = 1.0    # 1,2 s — firmware usually clears within 1-2 s
INTER_CHUNK_MARGIN = 0.1    # extra sleep after each chunk's audio duration
SLOW_FAIL_S        = 2.0    # if PlayStream blocks >N s then fails, the
                            # firmware almost certainly played the audio
                            # already — don't retry (we'd hear it twice).


def _read_wav_pcm(path: Path) -> Optional[bytes]:
    """Return the raw PCM bytes if the WAV matches the expected format, else None."""
    try:
        with wave.open(str(path), "rb") as wf:
            ch = wf.getnchannels()
            sw = wf.getsampwidth()
            fr = wf.getframerate()
            if ch != EXPECTED_CHANNELS or sw != EXPECTED_WIDTH or fr != EXPECTED_RATE:
                print(
                    f"[audio_player {_ts()}][WARN] {path}: expected "
                    f"{EXPECTED_RATE} Hz mono 16-bit; got "
                    f"{fr} Hz {ch}-ch {sw*8}-bit. Skipping.",
                    flush=True,
                )
                return None
            return wf.readframes(wf.getnframes())
    except Exception as e:
        print(f"[audio_player {_ts()}][WARN] failed to load {path}: {e}", flush=True)
        return None


class AudioPlayer:
    """Loads WAVs under ``assets/audio/<category>/<key>.wav`` and plays them on the G1."""

    def __init__(self, audio_client):
        self.audio_client = audio_client
        self._clips: Dict[Tuple[str, str], bytes] = {}
        # Set by callers via cancel() to interrupt an in-progress play when a
        # newer event arrives. play() checks the flag at every safe exit point
        # (before each PlayStream, inside retry back-offs, between chunks) and
        # bails out early so the worker thread can move to the next queue item.
        self._cancel = threading.Event()
        self._load_all()
        if self._clips:
            print(f"[audio_player {_ts()}] loaded {len(self._clips)} clip(s): "
                  f"{sorted(self._clips.keys())}", flush=True)
        else:
            print(f"[audio_player {_ts()}] no clips found under {AUDIO_ROOT}", flush=True)

    # ── library ─────────────────────────────────────────────────────────────
    def _load_all(self) -> None:
        if not AUDIO_ROOT.exists():
            return
        for category_dir in sorted(AUDIO_ROOT.iterdir()):
            if not category_dir.is_dir():
                continue
            for wav_path in sorted(category_dir.glob("*.wav")):
                pcm = _read_wav_pcm(wav_path)
                if pcm is not None:
                    self._clips[(category_dir.name, wav_path.stem)] = pcm

    def has(self, category: str, key: str) -> bool:
        return (category, key) in self._clips

    def cancel(self) -> None:
        """Signal the currently-playing clip (if any) to stop ASAP.

        Called when a newer announcement arrives — the worker thread should
        abandon the stale clip and move to the next one so audio stays in
        sync with the latest arm motion.

        IMPORTANT: we deliberately do NOT send AUDIO_STOP_PLAY to the
        firmware here. That DDS RPC blocks for up to ~10 s when the
        firmware is busy (e.g. during an arm replay), which is exactly
        when cancel() tends to be called. Setting the event flag is
        enough — play() checks it at every sleep/retry and returns False.
        The firmware will naturally finish draining the in-flight PCM
        buffer (a couple hundred ms of tail audio at most) while the
        worker thread moves on to the new queue item.
        """
        self._cancel.set()

    # ── playback ────────────────────────────────────────────────────────────
    def play(self, category: str, key: str) -> bool:
        """Blocking play. Returns True on success, False if clip missing / failed
        / cancelled. Clears the cancel flag at entry so each new play starts fresh.
        """
        self._cancel.clear()

        pcm = self._clips.get((category, key))
        if pcm is None:
            return False
        if self.audio_client is None:
            return False

        # NOTE: we used to call AUDIO_STOP_PLAY here and PlayStop below to be
        # defensive. Both are DDS RPCs that can *block 5-10 s* when the G1
        # firmware is busy (arm motion, recent TTS, etc.) — that's where most
        # of the "mystery lag" was coming from. Skipping them: a new sid per
        # play is enough for the firmware to treat this as a fresh stream.
        # cancel() also avoids the RPC; it only sets the _cancel flag.
        #
        # SetVolume stays fast (<50 ms), keep it to guard against firmware
        # resets between sessions.
        try:
            self.audio_client.SetVolume(100)
        except Exception:
            pass

        sid = f"saqr_{int(time.time() * 1000)}"
        offset = 0
        chunk0_attempts = 0
        chunkn_attempts = 0
        while offset < len(pcm):
            if self._cancel.is_set():
                print(f"[audio_player {_ts()}] cancelled mid-play (new announcement)",
                      flush=True)
                return False

            chunk = pcm[offset:offset + PLAY_CHUNK_BYTES]
            t_call_start = time.monotonic()
            code, _ = self.audio_client.PlayStream(PLAY_APP_NAME, sid, chunk)
            t_call = time.monotonic() - t_call_start

            if code != 0:
                # Slow-fail: firmware held the call for >SLOW_FAIL_S then
                # rejected. The audio almost certainly played during that
                # block — retrying would stutter / double-play. Treat as done.
                if t_call > SLOW_FAIL_S:
                    print(f"[audio_player {_ts()}][WARN] PlayStream rc={code} at byte "
                          f"{offset} after {t_call:.1f}s — assuming chunk played, "
                          f"skipping retry to avoid double audio", flush=True)
                    offset += len(chunk)
                    chunk_seconds = len(chunk) / (EXPECTED_RATE * EXPECTED_WIDTH)
                    if self._cancel.wait(timeout=chunk_seconds + INTER_CHUNK_MARGIN):
                        return False
                    continue

                # Fast-fail on chunk 0: retry with linear back-off.
                # new sid so firmware sees a fresh stream.
                if offset == 0 and chunk0_attempts < CHUNK0_RETRIES:
                    chunk0_attempts += 1
                    delay = CHUNK0_BACKOFF_S * chunk0_attempts
                    print(f"[audio_player {_ts()}][WARN] PlayStream rc={code} at byte 0 "
                          f"after {t_call:.2f}s; retry {chunk0_attempts}/{CHUNK0_RETRIES} "
                          f"in {delay:.1f}s", flush=True)
                    if self._cancel.wait(timeout=delay):
                        print(f"[audio_player {_ts()}] cancelled during retry back-off",
                              flush=True)
                        return False
                    sid = f"saqr_{int(time.time() * 1000)}"
                    continue

                # later chunks: firmware is still processing the previous chunk.
                # Retry a couple of times with back-off, keeping the same sid.
                if offset > 0 and chunkn_attempts < CHUNKN_RETRIES:
                    chunkn_attempts += 1
                    delay = CHUNKN_BACKOFF_S * chunkn_attempts
                    print(f"[audio_player {_ts()}][WARN] PlayStream rc={code} at byte {offset}; "
                          f"mid-stream retry {chunkn_attempts}/{CHUNKN_RETRIES} in {delay:.1f}s",
                          flush=True)
                    if self._cancel.wait(timeout=delay):
                        print(f"[audio_player {_ts()}] cancelled during mid-stream retry",
                              flush=True)
                        return False
                    continue

                print(f"[audio_player {_ts()}][WARN] PlayStream rc={code} at byte {offset} "
                      f"(retries exhausted)", flush=True)
                return False

            if offset == 0 and chunk0_attempts > 0:
                print(f"[audio_player {_ts()}] chunk 0 succeeded after "
                      f"{chunk0_attempts} retry/retries", flush=True)
            elif offset > 0 and chunkn_attempts > 0:
                print(f"[audio_player {_ts()}] chunk at byte {offset} succeeded after "
                      f"{chunkn_attempts} retry/retries", flush=True)
                chunkn_attempts = 0   # reset for any subsequent chunk

            offset += len(chunk)
            # Wait for the chunk to finish playing before sending the next.
            chunk_seconds = len(chunk) / (EXPECTED_RATE * EXPECTED_WIDTH)
            if self._cancel.wait(timeout=chunk_seconds + INTER_CHUNK_MARGIN):
                print(f"[audio_player {_ts()}] cancelled between chunks", flush=True)
                return False

        # Don't call PlayStop — see the note at the top of this method.
        # The firmware will finish draining the last chunk's PCM buffer on
        # its own once we stop sending more data.
        return True