#!/usr/bin/env python3
"""End-to-end sync tests through the Marionette stack.

Tests audio-motion synchronization by playing a synthetic move through
Marionette's full playback pipeline (push_audio_sample + motion loop)
and measuring beep-to-collision intervals with a laptop microphone.

Requires:
  - Marionette running on the robot (via deploy_wireless.sh or daemon)
  - Laptop microphone connected and working
  - Robot accessible at --host (default: reachy-mini.local)

Usage:
    python tests/test_marionette_sync.py [--host reachy-mini.local]
"""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
import time
from pathlib import Path

import numpy as np
import requests
import sounddevice as sd
import soundfile as sf

sys.path.insert(0, str(Path(__file__).parent))
from audio_analysis import detect_beep_onsets, detect_transient_onsets

# ── Timing (same as all previous sync tests) ──────────────────────────
BEEP_TIMES = [1.0, 2.3, 4.0, 6.3, 9.4]
BEEP_COLLISION_OFFSET = 1.0
COLLISION_TIMES = [t + BEEP_COLLISION_OFFSET for t in BEEP_TIMES]

# Audio
BEEP_FREQ = 2000.0
BEEP_DURATION = 0.2
BEEP_AMPLITUDE = 0.9
ROBOT_SR = 16000

# Collision
RIGHT_REST = -0.68
LEFT_REST = 0.0
LEFT_COLLISION = 0.70
HOLD_DURATION = 0.2

MOTION_SR = 100
MOVE_ID = "sync-test-e2e"

ROBOT_USER = "pollen"
LAPTOP_SR = 48000
MIC_DURATION = 35.0  # longer: Marionette has goto + playback
MARIONETTE_PORT = 8042


def get_marionette_url(host: str) -> str:
    return f"http://{host}:{MARIONETTE_PORT}"


def wait_for_mode(base_url: str, target: str, timeout: float = 30.0) -> dict:
    """Poll GET /api/state until mode matches target."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            r = requests.get(f"{base_url}/api/state", timeout=3)
            state = r.json()
            if state.get("mode") == target:
                return state
        except Exception:
            pass
        time.sleep(0.3)
    raise TimeoutError(f"Mode never reached '{target}' within {timeout}s")


def ensure_idle(base_url: str) -> dict:
    """Make sure Marionette is idle, stopping anything in progress."""
    try:
        state = requests.get(f"{base_url}/api/state", timeout=3).json()
    except Exception as exc:
        raise RuntimeError(f"Cannot reach Marionette at {base_url}: {exc}") from exc

    mode = state.get("mode", "unknown")
    if mode == "idle":
        return state

    # Try to stop whatever is running
    if mode == "playing":
        requests.post(f"{base_url}/api/play/stop", timeout=3)
    elif mode in {"recording", "countdown", "preparing"}:
        requests.post(f"{base_url}/api/record/stop", timeout=3)
    return wait_for_mode(base_url, "idle", timeout=10)


def generate_move_json() -> dict:
    """Generate Marionette-format move data with collision trajectory."""
    total_duration = max(COLLISION_TIMES) + HOLD_DURATION + 1.0
    dt = 1.0 / MOTION_SR
    n_frames = int(total_duration * MOTION_SR)

    identity_head = np.eye(4).tolist()
    left_targets = np.full(n_frames, LEFT_REST, dtype=np.float64)
    for ct in COLLISION_TIMES:
        start = int(ct * MOTION_SR)
        end = min(int((ct + HOLD_DURATION) * MOTION_SR), n_frames)
        left_targets[start:end] = LEFT_COLLISION

    timestamps = []
    frames = []
    for i in range(n_frames):
        t = round(i * dt, 4)
        timestamps.append(t)
        frames.append({
            "head": identity_head,
            "antennas": [float(left_targets[i]), RIGHT_REST],
            "body_yaw": 0.0,
            "check_collision": False,
        })

    return {
        "description": "E2E sync test: beeps + antenna collisions",
        "time": timestamps,
        "set_target_data": frames,
    }


def generate_move_wav(path: Path) -> None:
    """Generate WAV with beeps at known times."""
    total_duration = max(COLLISION_TIMES) + HOLD_DURATION + 1.0
    n_audio = int(total_duration * ROBOT_SR)
    audio = np.zeros(n_audio, dtype=np.float32)
    for bt in BEEP_TIMES:
        start = int(bt * ROBOT_SR)
        n_beep = int(BEEP_DURATION * ROBOT_SR)
        if start + n_beep > n_audio:
            continue
        t_arr = np.arange(n_beep, dtype=np.float32) / ROBOT_SR
        beep = BEEP_AMPLITUDE * np.sin(2 * np.pi * BEEP_FREQ * t_arr).astype(np.float32)
        fade = int(0.005 * ROBOT_SR)
        if fade > 0 and 2 * fade < n_beep:
            beep[:fade] *= np.linspace(0, 1, fade, dtype=np.float32)
            beep[-fade:] *= np.linspace(1, 0, fade, dtype=np.float32)
        audio[start:start + n_beep] += beep
    sf.write(str(path), audio, ROBOT_SR)


def inject_move(host: str, base_url: str) -> None:
    """Inject the synthetic move into the running Marionette's dataset."""
    # Get the active dataset path from Marionette
    state = requests.get(f"{base_url}/api/state", timeout=3).json()
    dataset_path = state["config"]["active_dataset_path"]
    active_id = state.get("datasets", {}).get("active_id")
    print(f"  Active dataset path: {dataset_path}")

    # Generate move files locally
    import tempfile
    with tempfile.TemporaryDirectory() as tmpdir:
        tmpdir = Path(tmpdir)
        json_path = tmpdir / f"{MOVE_ID}.json"
        wav_path = tmpdir / f"{MOVE_ID}.wav"

        move_data = generate_move_json()
        json_path.write_text(json.dumps(move_data), encoding="utf-8")
        generate_move_wav(wav_path)

        total_duration = max(COLLISION_TIMES) + HOLD_DURATION + 1.0
        print(f"  Generated move: {total_duration:.1f}s, {len(BEEP_TIMES)} beeps, {len(COLLISION_TIMES)} collisions")

        # SCP to robot's dataset directory
        for local_file in [json_path, wav_path]:
            remote_target = f"{ROBOT_USER}@{host}:{dataset_path}/{local_file.name}"
            result = subprocess.run(
                ["scp", "-o", "ConnectTimeout=5", str(local_file), remote_target],
                capture_output=True, text=True, timeout=15,
            )
            if result.returncode != 0:
                raise RuntimeError(f"SCP failed: {result.stderr}")

    # Trigger Marionette to re-scan recordings by re-selecting the active dataset
    if active_id:
        requests.post(
            f"{base_url}/api/datasets/select",
            json={"dataset_id": active_id},
            timeout=5,
        )
    time.sleep(0.5)

    # Verify the move is visible
    state = requests.get(f"{base_url}/api/state", timeout=3).json()
    move_ids = [m["id"] for m in state.get("moves", [])]
    if MOVE_ID not in move_ids:
        raise RuntimeError(
            f"Move '{MOVE_ID}' not found after injection. "
            f"Available: {move_ids}"
        )
    print(f"  Move '{MOVE_ID}' injected and visible in Marionette")


def cleanup_move(host: str, base_url: str) -> None:
    """Remove the synthetic test move."""
    try:
        requests.delete(f"{base_url}/api/moves/{MOVE_ID}", timeout=5)
    except Exception:
        pass


def plot_results(
    mic_audio, mic_sr, detected_beeps, detected_collisions, pairs, output_path,
):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    mic_t = np.arange(len(mic_audio)) / mic_sr

    fig, ax = plt.subplots(1, 1, figsize=(18, 6))
    ax.plot(mic_t, mic_audio, "k-", linewidth=0.3, alpha=0.5)
    ax.set_ylabel("Mic amplitude")
    ax.set_title("Marionette E2E Playback Sync Test — Laptop Mic Recording")
    ax.grid(True, alpha=0.3)

    for i, bt in enumerate(detected_beeps):
        ax.axvline(bt, color="blue", linestyle="-", linewidth=1.2, alpha=0.7,
                    label="Detected beep" if i == 0 else None)
    for i, ct in enumerate(detected_collisions):
        ax.axvline(ct, color="red", linestyle="-", linewidth=1.2, alpha=0.7,
                    label="Detected collision" if i == 0 else None)
    for p in pairs:
        mid = (p["beep_t"] + p["collision_t"]) / 2
        ax.annotate(f'{p["interval_ms"]:.0f}ms', xy=(mid, 0),
                    ha="center", fontsize=9, color="purple", fontweight="bold",
                    bbox=dict(boxstyle="round,pad=0.2", facecolor="lightyellow", alpha=0.8))

    ax.legend(loc="upper right", fontsize=9)

    all_events = detected_beeps + detected_collisions
    if all_events:
        ax.set_xlim(min(all_events) - 1.0, max(all_events) + 1.0)
    ax.set_xlabel("Time since mic start (s)")

    fig.tight_layout()
    fig.savefig(str(output_path), dpi=150)
    plt.close(fig)
    print(f"  Plot saved to {output_path}")


def run_playback_test(host: str) -> dict:
    """Run the Marionette playback sync test.

    Returns a results dict with pairs, errors, and verdict.
    """
    base_url = get_marionette_url(host)
    print(f"\n{'='*60}")
    print("Marionette E2E Playback Sync Test")
    print(f"{'='*60}")
    print(f"  Server: {base_url}")
    print(f"  Beep times:      {BEEP_TIMES}")
    print(f"  Collision times:  {COLLISION_TIMES}")
    print(f"  Expected interval: {BEEP_COLLISION_OFFSET*1000:.0f}ms\n")

    # Step 1: Ensure idle
    print("[1/6] Ensuring Marionette is idle...")
    ensure_idle(base_url)
    print("  Idle.")

    # Step 2: Inject synthetic move
    print("[2/6] Injecting synthetic move into dataset...")
    inject_move(host, base_url)

    # Step 3: Start mic recording
    print(f"\n[3/6] Starting mic recording ({MIC_DURATION}s)...")
    mic_start = time.monotonic()
    mic_data = sd.rec(
        int(MIC_DURATION * LAPTOP_SR),
        samplerate=LAPTOP_SR, channels=1, dtype="float32",
    )

    # Step 4: Trigger playback
    time.sleep(0.3)
    print("[4/6] Triggering playback via API...")
    play_start = time.monotonic()
    r = requests.post(f"{base_url}/api/play", json={"move_id": MOVE_ID}, timeout=5)
    if r.status_code != 200:
        sd.stop()
        raise RuntimeError(f"Play failed: {r.status_code} {r.text}")
    print(f"  Play accepted at mic_t={play_start - mic_start:.3f}s")

    # Step 5: Wait for playback to finish
    print("[5/6] Waiting for playback to finish...")
    try:
        state = wait_for_mode(base_url, "idle", timeout=30)
        play_end = time.monotonic()
        print(f"  Playback done at mic_t={play_end - mic_start:.3f}s")
    except TimeoutError:
        print("  WARNING: Playback did not finish in time")

    sd.wait()
    captured = mic_data.flatten()
    print(f"  Mic recording done ({len(captured)/LAPTOP_SR:.1f}s)")

    mic_path = Path("tests/marionette_sync_mic.wav")
    sf.write(str(mic_path), captured, LAPTOP_SR)
    print(f"  Saved to {mic_path}")

    # Step 6: Analyze
    print(f"\n[6/6] Analyzing...")
    detected_beeps = detect_beep_onsets(
        captured, LAPTOP_SR, freq=BEEP_FREQ, bandwidth=150.0,
        threshold_db=-12.0, min_separation=1.0,
    )
    detected_collisions = detect_transient_onsets(
        captured, LAPTOP_SR, highpass_freq=3000.0,
    )
    print(f"  Detected {len(detected_beeps)} beeps at: {[f'{t:.3f}' for t in detected_beeps]}")
    print(f"  Detected {len(detected_collisions)} collisions at: {[f'{t:.3f}' for t in detected_collisions]}")

    # Match pairs
    print(f"\n{'='*60}")
    print("Beep → Collision Interval Analysis (Marionette E2E)")
    print(f"  (Expected interval: {BEEP_COLLISION_OFFSET*1000:.0f}ms)")
    print(f"{'='*60}")

    pairs = []
    for i, bt in enumerate(detected_beeps):
        candidates = [ct for ct in detected_collisions if 0.3 < (ct - bt) < 2.0]
        if not candidates:
            print(f"  Beep {i+1} at {bt:.3f}s: NO COLLISION FOUND")
            continue
        nearest = min(candidates, key=lambda ct: abs((ct - bt) - BEEP_COLLISION_OFFSET))
        interval_ms = (nearest - bt) * 1000
        error_ms = interval_ms - BEEP_COLLISION_OFFSET * 1000
        pairs.append({
            "beep_t": bt,
            "collision_t": nearest,
            "interval_ms": interval_ms,
            "error_ms": error_ms,
        })
        print(f"  Pair {len(pairs)}: beep {bt:.3f}s → collision {nearest:.3f}s = "
              f"{interval_ms:.0f}ms (error {error_ms:+.0f}ms)")

    result = {
        "test": "marionette_playback_sync",
        "n_beeps_detected": len(detected_beeps),
        "n_collisions_detected": len(detected_collisions),
        "n_pairs": len(pairs),
        "pairs": pairs,
    }

    if pairs:
        errors = [p["error_ms"] for p in pairs]
        intervals = [p["interval_ms"] for p in pairs]
        result["mean_interval_ms"] = float(np.mean(intervals))
        result["mean_error_ms"] = float(np.mean(errors))
        result["std_error_ms"] = float(np.std(errors))
        result["min_error_ms"] = float(min(errors))
        result["max_error_ms"] = float(max(errors))
        print(f"\n  Pairs matched: {len(pairs)}/{len(BEEP_TIMES)}")
        print(f"  Mean interval: {np.mean(intervals):.0f}ms (expected {BEEP_COLLISION_OFFSET*1000:.0f}ms)")
        print(f"  Mean error:    {np.mean(errors):+.0f}ms")
        print(f"  Std error:     {np.std(errors):.0f}ms")
        print(f"  Min/Max error: {min(errors):+.0f}ms / {max(errors):+.0f}ms")

    # Plot
    plot_path = Path("tests/marionette_sync_plot.png")
    plot_results(captured, LAPTOP_SR, detected_beeps, detected_collisions, pairs, plot_path)

    # Cleanup
    cleanup_move(host, base_url)

    success = len(pairs) >= len(BEEP_TIMES) - 1
    result["success"] = success

    print(f"\n{'='*60}")
    if success:
        print("RESULT: PASS — Beep-collision pairs detected via Marionette E2E")
    else:
        print("RESULT: FAIL — Could not reliably detect pairs")
    print(f"{'='*60}\n")

    return result


def run_recording_test(host: str) -> dict:
    """Test the 3-2-1 countdown timing accuracy.

    Triggers a recording with a known audio file, records with the
    laptop mic, and measures:
      - Countdown beep timing (3 beeps at 440Hz, 1s apart)
      - "Go" beep timing (880Hz)
      - When the uploaded audio actually starts playing
    """
    base_url = get_marionette_url(host)
    print(f"\n{'='*60}")
    print("Marionette Recording Countdown Sync Test")
    print(f"{'='*60}")
    print(f"  Server: {base_url}\n")

    # Step 1: Ensure idle
    print("[1/5] Ensuring Marionette is idle...")
    ensure_idle(base_url)

    # Step 2: Upload a known audio file (a single 2kHz beep at t=0.5)
    # This beep will play after the countdown, so we can measure
    # the delay from "go" beep to actual audio playback start.
    print("[2/5] Generating and uploading test audio...")
    import tempfile
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        test_wav = Path(f.name)

    # 3-second audio with a single 2kHz marker beep at t=0.2
    # (early, so it's clearly after the go beep)
    marker_time = 0.2
    audio_duration = 3.0
    n_audio = int(audio_duration * 48000)  # 48kHz for upload
    audio = np.zeros(n_audio, dtype=np.float32)
    start = int(marker_time * 48000)
    n_beep = int(BEEP_DURATION * 48000)
    t_arr = np.arange(n_beep, dtype=np.float32) / 48000
    beep = BEEP_AMPLITUDE * np.sin(2 * np.pi * BEEP_FREQ * t_arr).astype(np.float32)
    fade = int(0.005 * 48000)
    if fade > 0:
        beep[:fade] *= np.linspace(0, 1, fade, dtype=np.float32)
        beep[-fade:] *= np.linspace(1, 0, fade, dtype=np.float32)
    audio[start:start + n_beep] = beep
    sf.write(str(test_wav), audio, 48000)
    print(f"  Test audio: {audio_duration}s, marker beep at {marker_time}s (2kHz)")

    # Upload via Marionette API
    with open(test_wav, "rb") as f:
        r = requests.post(
            f"{base_url}/api/upload-audio",
            files={"file": ("countdown_test.wav", f, "audio/wav")},
            timeout=10,
        )
    test_wav.unlink()
    if r.status_code != 200:
        raise RuntimeError(f"Upload failed: {r.status_code} {r.text}")
    upload_info = r.json()
    upload_id = upload_info["upload_id"]
    print(f"  Uploaded: id={upload_id}")

    # Step 3: Start mic recording
    print(f"\n[3/5] Starting mic recording (15s)...")
    mic_duration = 15.0
    mic_start = time.monotonic()
    mic_data = sd.rec(
        int(mic_duration * LAPTOP_SR),
        samplerate=LAPTOP_SR, channels=1, dtype="float32",
    )

    # Step 4: Trigger recording
    time.sleep(0.3)
    print("[4/5] Triggering recording with audio...")
    rec_start = time.monotonic()
    r = requests.post(
        f"{base_url}/api/record",
        json={
            "label": "countdown-test",
            "duration": 3.0,
            "record_audio": False,
            "record_motion": True,
            "uploaded_audio_id": upload_id,
        },
        timeout=5,
    )
    if r.status_code != 200:
        sd.stop()
        raise RuntimeError(f"Record failed: {r.status_code} {r.text}")
    rec_move_id = r.json().get("move_id")
    print(f"  Recording accepted at mic_t={rec_start - mic_start:.3f}s, move_id={rec_move_id}")

    # Wait for recording to finish (countdown ~3s + recording ~3s)
    print("[5/5] Waiting for recording to finish...")
    try:
        wait_for_mode(base_url, "idle", timeout=20)
        rec_end = time.monotonic()
        print(f"  Recording done at mic_t={rec_end - mic_start:.3f}s")
    except TimeoutError:
        print("  WARNING: Recording did not finish in time")

    sd.wait()
    captured = mic_data.flatten()

    mic_path = Path("tests/marionette_countdown_mic.wav")
    sf.write(str(mic_path), captured, LAPTOP_SR)
    print(f"  Saved to {mic_path}")

    # Analyze: detect countdown beeps (440Hz) and go beep (880Hz) and marker (2kHz)
    print(f"\n{'='*60}")
    print("Countdown Timing Analysis")
    print(f"{'='*60}")

    # Detect countdown beeps at 440Hz
    countdown_beeps = detect_beep_onsets(
        captured, LAPTOP_SR, freq=440.0, bandwidth=80.0,
        threshold_db=-12.0, min_separation=0.8,
    )
    print(f"  Countdown beeps (440Hz): {len(countdown_beeps)} at {[f'{t:.3f}' for t in countdown_beeps]}")

    # Detect "go" beep at 880Hz
    go_beeps = detect_beep_onsets(
        captured, LAPTOP_SR, freq=880.0, bandwidth=80.0,
        threshold_db=-12.0, min_separation=0.5,
    )
    print(f"  Go beep (880Hz): {len(go_beeps)} at {[f'{t:.3f}' for t in go_beeps]}")

    # Detect marker beep at 2kHz (from uploaded audio)
    marker_beeps = detect_beep_onsets(
        captured, LAPTOP_SR, freq=BEEP_FREQ, bandwidth=150.0,
        threshold_db=-12.0, min_separation=0.5,
    )
    print(f"  Marker beep (2kHz): {len(marker_beeps)} at {[f'{t:.3f}' for t in marker_beeps]}")

    result = {
        "test": "marionette_recording_countdown",
        "countdown_beeps": countdown_beeps,
        "go_beeps": go_beeps,
        "marker_beeps": marker_beeps,
    }

    # Analyze countdown spacing
    if len(countdown_beeps) >= 2:
        gaps = [countdown_beeps[i+1] - countdown_beeps[i] for i in range(len(countdown_beeps)-1)]
        print(f"\n  Countdown gaps: {[f'{g:.3f}s' for g in gaps]} (expected ~1.0s each)")
        result["countdown_gaps"] = gaps
        gap_errors = [abs(g - 1.0) * 1000 for g in gaps]
        print(f"  Gap errors: {[f'{e:.0f}ms' for e in gap_errors]}")

    # Analyze go-to-marker delay
    # Filter go beeps: discard any that overlap with countdown beeps (440Hz
    # harmonic leaks into the 880Hz band).  The real "go" beep comes AFTER the
    # last countdown beep.
    if countdown_beeps and go_beeps:
        last_cd = countdown_beeps[-1]
        go_beeps_filtered = [t for t in go_beeps if t > last_cd + 0.3]
        print(f"  Go beeps after countdown: {[f'{t:.3f}' for t in go_beeps_filtered]}")
    else:
        go_beeps_filtered = go_beeps

    if go_beeps_filtered and marker_beeps:
        go_t = go_beeps_filtered[0]
        marker_t = marker_beeps[0]
        delay_ms = (marker_t - go_t) * 1000
        print(f"\n  Go beep → Marker beep: {delay_ms:.0f}ms")
        print(f"  (Marker is at {marker_time}s in audio file)")
        print(f"  Expected if no pipeline latency: ~{marker_time*1000:.0f}ms")
        print(f"  Extra delay (pipeline latency): ~{delay_ms - marker_time*1000:.0f}ms")
        result["go_to_marker_ms"] = delay_ms
        result["estimated_pipeline_latency_ms"] = delay_ms - marker_time * 1000

    # Cleanup: delete the test recording
    if rec_move_id:
        try:
            requests.delete(f"{base_url}/api/moves/{rec_move_id}", timeout=5)
        except Exception:
            pass

    success = len(countdown_beeps) >= 2 and len(marker_beeps) >= 1
    result["success"] = success

    print(f"\n{'='*60}")
    if success:
        print("RESULT: PASS — Countdown and marker beeps detected")
    else:
        print("RESULT: FAIL — Could not detect expected beeps")
    print(f"{'='*60}\n")

    return result


def main():
    parser = argparse.ArgumentParser(description="Marionette E2E sync tests")
    parser.add_argument("--host", default="reachy-mini.local")
    parser.add_argument("--test", choices=["playback", "recording", "both"],
                        default="both", help="Which test to run")
    args = parser.parse_args()

    results = {}

    if args.test in ("playback", "both"):
        results["playback"] = run_playback_test(args.host)

    if args.test in ("recording", "both"):
        results["recording"] = run_recording_test(args.host)

    # Save results
    results_path = Path("tests/marionette_sync_results.json")
    results_path.write_text(json.dumps(results, indent=2, default=str))
    print(f"\nResults saved to {results_path}")

    # Summary
    print(f"\n{'='*60}")
    print("SUMMARY")
    print(f"{'='*60}")
    for name, r in results.items():
        status = "PASS" if r.get("success") else "FAIL"
        if name == "playback" and "mean_error_ms" in r:
            print(f"  {name}: {status} — mean error {r['mean_error_ms']:+.0f}ms, "
                  f"std {r['std_error_ms']:.0f}ms")
        elif name == "recording" and "go_to_marker_ms" in r:
            print(f"  {name}: {status} — go→marker {r['go_to_marker_ms']:.0f}ms, "
                  f"pipeline latency ~{r.get('estimated_pipeline_latency_ms', 0):.0f}ms")
        else:
            print(f"  {name}: {status}")
    print(f"{'='*60}\n")

    all_pass = all(r.get("success", False) for r in results.values())
    return 0 if all_pass else 1


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