#!/opt/homebrew/bin/python3.11
"""
Speaker Audio Player - 說話人語音播放器
從 ASRX 結果中提取並播放每個說話人的語音片段
"""

import json
import argparse
import subprocess
import tempfile
import os
from pathlib import Path
from typing import List, Dict


def load_asrx_result(result_path: str) -> Dict:
    """載入 ASRX 結果"""
    with open(result_path, "r", encoding="utf-8") as f:
        return json.load(f)


def extract_audio_segment(
    audio_path: str, start_sec: float, end_sec: float, output_path: str
) -> bool:
    """
    使用 ffmpeg 提取音頻片段

    Args:
        audio_path: 原始音頻路徑
        start_sec: 開始時間（秒）
        end_sec: 結束時間（秒）
        output_path: 輸出路徑

    Returns:
        bool: 是否成功
    """
    duration = end_sec - start_sec

    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        audio_path,
        "-ss",
        str(start_sec),
        "-t",
        str(duration),
        "-acodec",
        "pcm_s16le",
        "-ar",
        "16000",
        "-ac",
        "1",
        output_path,
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True)
        return result.returncode == 0
    except Exception as e:
        print(f"Error extracting audio: {e}")
        return False


def play_audio(audio_path: str) -> bool:
    """
    播放音頻文件

    使用 macOS 的 afplay 或 Linux 的 aplay
    """
    try:
        # 嘗試使用 afplay (macOS)
        if os.path.exists("/usr/bin/afplay"):
            subprocess.run(["afplay", audio_path], check=True)
        # 嘗試使用 aplay (Linux)
        elif os.path.exists("/usr/bin/aplay"):
            subprocess.run(["aplay", audio_path], check=True)
        else:
            print(
                "No audio player found. Please install afplay (macOS) or aplay (Linux)"
            )
            return False
        return True
    except Exception as e:
        print(f"Error playing audio: {e}")
        return False


def group_segments_by_speaker(segments: List[Dict]) -> Dict[str, List[Dict]]:
    """將語音片段按說話人分組"""
    speaker_segments = {}

    for seg in segments:
        speaker = seg["speaker"]
        if speaker not in speaker_segments:
            speaker_segments[speaker] = []
        speaker_segments[speaker].append(seg)

    # 按開始時間排序
    for speaker in speaker_segments:
        speaker_segments[speaker].sort(key=lambda x: x["start"])

    return speaker_segments


def play_speaker_segments(
    audio_path: str,
    result_path: str,
    speaker_id: str = None,
    limit: int = None,
    temp_dir: str = None,
):
    """
    播放指定說話人的語音片段

    Args:
        audio_path: 原始音頻路徑
        result_path: ASRX 結果 JSON 路徑
        speaker_id: 說話人 ID（None=播放所有）
        limit: 最多播放幾個片段（None=全部）
        temp_dir: 臨時目錄
    """
    # 載入結果
    print(f"[Load] Loading ASRX result: {result_path}")
    result = load_asrx_result(result_path)

    segments = result.get("segments", [])
    total_duration = result.get("total_duration", 0)

    print(f"[Info] Total segments: {len(segments)}")
    print(f"[Info] Total duration: {total_duration / 60:.1f} minutes")

    # 分組
    speaker_segments = group_segments_by_speaker(segments)

    # 選擇說話人
    if speaker_id:
        speakers_to_play = [speaker_id]
    else:
        speakers_to_play = sorted(speaker_segments.keys())

    # 創建臨時目錄
    if temp_dir is None:
        temp_dir = tempfile.mkdtemp(prefix="speaker_audio_")

    print(f"\n[Info] Temp directory: {temp_dir}")
    print(f"[Info] Speakers to play: {speakers_to_play}")
    print("=" * 60)

    # 播放每個說話人的片段
    for speaker in speakers_to_play:
        if speaker not in speaker_segments:
            print(f"\n[Warning] Speaker {speaker} not found!")
            continue

        segs = speaker_segments[speaker]
        if limit:
            segs = segs[:limit]

        print(f"\n▶️  {speaker} ({len(segs)} segments)")
        print("-" * 60)

        for i, seg in enumerate(segs, 1):
            start = seg["start"]
            end = seg["end"]
            duration = seg["duration"]

            # 提取音頻
            temp_audio = os.path.join(temp_dir, f"{speaker}_{i:03d}.wav")

            print(
                f"  [{i:3d}] {start:7.2f}s - {end:7.2f}s ({duration:5.2f}s) ... ",
                end="",
                flush=True,
            )

            if extract_audio_segment(audio_path, start, end, temp_audio):
                print("✅", end="", flush=True)

                # 播放
                if play_audio(temp_audio):
                    print(" ▶️  Played")
                else:
                    print(" ❌ Play failed")
            else:
                print(" ❌ Extract failed")

        print()


def show_speaker_stats(result_path: str):
    """顯示說話人統計資訊"""
    result = load_asrx_result(result_path)

    segments = result.get("segments", [])
    speaker_segments = group_segments_by_speaker(segments)

    print("\n" + "=" * 60)
    print("說話人統計")
    print("=" * 60)

    # 按時長排序
    speaker_stats = []
    for speaker, segs in speaker_segments.items():
        total_duration = sum(seg["duration"] for seg in segs)
        speaker_stats.append((speaker, len(segs), total_duration))

    speaker_stats.sort(key=lambda x: x[2], reverse=True)

    total_duration = result.get("total_duration", 0)

    for speaker, count, duration in speaker_stats:
        pct = duration / total_duration * 100 if total_duration > 0 else 0
        print(f"{speaker:12} {count:4} segments  {duration:8.1f}s  ({pct:5.1f}%)")

    print("=" * 60)


def main():
    parser = argparse.ArgumentParser(
        description="Speaker Audio Player - 播放說話人語音片段",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # 顯示說話人統計
  python3 speaker_audio_player.py --stats result.json
  
  # 播放所有說話人的前 3 個片段
  python3 speaker_audio_player.py audio.wav result.json --limit 3
  
  # 播放特定說話人的所有片段
  python3 speaker_audio_player.py audio.wav result.json --speaker SPEAKER_0
  
  # 播放 SPEAKER_1 的前 5 個片段
  python3 speaker_audio_player.py audio.wav result.json --speaker SPEAKER_1 --limit 5
        """,
    )

    parser.add_argument("audio_path", nargs="?", help="原始音頻文件路徑")
    parser.add_argument("result_path", help="ASRX 結果 JSON 路徑")
    parser.add_argument("--stats", action="store_true", help="只显示說話人統計")
    parser.add_argument("--speaker", type=str, help="指定說話人 ID（如 SPEAKER_0）")
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="每個說話人最多播放幾個片段（None=全部）",
    )
    parser.add_argument("--temp-dir", type=str, default=None, help="臨時目錄路徑")

    args = parser.parse_args()

    if args.stats:
        show_speaker_stats(args.result_path)
        return

    if not args.audio_path:
        print("Error: audio_path is required unless --stats is specified")
        parser.print_help()
        return

    if not Path(args.audio_path).exists():
        print(f"Error: Audio file not found: {args.audio_path}")
        return

    if not Path(args.result_path).exists():
        print(f"Error: Result file not found: {args.result_path}")
        return

    play_speaker_segments(
        args.audio_path,
        args.result_path,
        speaker_id=args.speaker,
        limit=args.limit,
        temp_dir=args.temp_dir,
    )


if __name__ == "__main__":
    main()
