#!/opt/homebrew/bin/python3.11
"""
pyannote.audio 測試腳本
測試說話人分離功能
"""

import sys
import json

def test_pyannote(audio_path, output_path):
    """測試 pyannote.audio 說話人分離"""
    
    print(f"[pyannote] Testing on: {audio_path}")
    
    try:
        from pyannote.audio import Pipeline
        
        # 載入模型
        print("[pyannote] Loading model...")
        pipeline = Pipeline.from_pretrained(
            "pyannote/speaker-diarization-3.1"
        )
        
        print("[pyannote] Model loaded successfully")
        
        # 執行說話人分離
        print("[pyannote] Performing speaker diarization...")
        diarization = pipeline(audio_path)
        
        # 收集結果
        segments = []
        for turn, _, speaker in diarization.itertracks(yield_label=True):
            segments.append({
                "start": turn.start,
                "end": turn.end,
                "speaker": speaker
            })
        
        # 輸出結果
        result = {
            "audio": audio_path,
            "num_speakers": len(set(s["speaker"] for s in segments)),
            "segments": segments
        }
        
        with open(output_path, "w") as f:
            json.dump(result, f, indent=2)
        
        print("\n=== 測試結果 ===")
        print(f"說話人數量：{result['num_speakers']}")
        print(f"片段數量：{len(segments)}")
        print(f"輸出檔案：{output_path}")
        
        # 顯示前 5 段
        print("\n前 5 段：")
        for i, seg in enumerate(segments[:5], 1):
            print(f"  {i}. [{seg['start']:.2f}s - {seg['end']:.2f}s] {seg['speaker']}")
        
        return True
        
    except Exception as e:
        print(f"[pyannote] Error: {e}")
        import traceback
        traceback.print_exc()
        
        # 寫出錯誤資訊
        result = {
            "audio": audio_path,
            "error": str(e)
        }
        with open(output_path, "w") as f:
            json.dump(result, f, indent=2)
        
        return False


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python3 test_pyannote_audio.py <audio_path> <output_path>")
        sys.exit(1)
    
    audio_path = sys.argv[1]
    output_path = sys.argv[2]
    
    success = test_pyannote(audio_path, output_path)
    sys.exit(0 if success else 1)