#!/opt/homebrew/bin/python3.11
"""
Audio Track Probe - Audio track detection and VAD classification

Used during S0 Register phase to classify audio tracks:
- no_audio: No audio track
- silent_audio: Audio track but no speech detected
- music_only: Audio with no speech (music/sound effects)
- speech_only: Audio with speech only
- speech_with_music: Audio with speech and background music

Usage:
    python audio_track_probe.py --file /path/to/video.mp4
    python audio_track_probe.py --file /path/to/video.mp4 --json

Output (text):
    music_only
    
Output (JSON):
    {"classification": "music_only", "speech_ratio": 0.0, "speech_segments": 0, "duration": 93.3}
"""

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

try:
    import torch
    import numpy as np
    from scipy.io import wavfile
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False


def get_audio_tracks(file_path: str) -> list[dict]:
    """
    Get audio track information using ffprobe.
    
    Returns:
        List of audio track dicts with: index, codec, channels, language, title
    """
    cmd = [
        "ffprobe", "-v", "quiet",
        "-print_format", "json",
        "-show_streams",
        "-select_streams", "a",
        file_path
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        return []
    
    data = json.loads(result.stdout)
    streams = data.get("streams", [])
    
    tracks = []
    for s in streams:
        track = {
            "index": s.get("index", 0),
            "codec": s.get("codec_name", "unknown"),
            "channels": s.get("channels", 2),
            "language": s.get("tags", {}).get("language", ""),
            "title": s.get("tags", {}).get("title", ""),
        }
        tracks.append(track)
    
    return tracks


def select_best_track(tracks: list[dict]) -> int | None:
    """
    Select the best audio track for VAD analysis.
    
    Priority (原聲優先):
    1. Language = original/und/unknown (assumed original)
    2. Language matches common original track codes
    3. Most channels
    4. First track
    
    Returns:
        Stream index of best track, or None if no tracks
    """
    if not tracks:
        return None
    
    # Priority 1: original/und/unknown language
    for t in tracks:
        lang = t.get("language", "").lower()
        if lang in ("", "und", "original", "unknown"):
            return t["index"]
    
    # Priority 2: common original track languages
    original_langs = ("zho", "chi", "jpn", "jap", "kor", "tha", "vie")
    for t in tracks:
        lang = t.get("language", "").lower()
        if lang in original_langs:
            return t["index"]
    
    # Priority 3: Most channels
    tracks_sorted = sorted(tracks, key=lambda x: x.get("channels", 0), reverse=True)
    return tracks_sorted[0]["index"]


def extract_audio_for_vad(file_path: str, stream_index: int | None = None) -> str | None:
    """
    Extract audio to temp WAV file for VAD analysis.
    
    Returns:
        Path to temp WAV file, or None if extraction failed
    """
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        output_path = f.name
    
    cmd = ["ffmpeg", "-y", "-v", "quiet"]
    
    if stream_index is not None:
        cmd.extend(["-stream_loop", "1", "-i", file_path, "-map", f"0:{stream_index}"])
    else:
        cmd.extend(["-i", file_path])
    
    cmd.extend([
        "-vn", "-ac", "1", "-ar", "16000",
        "-acodec", "pcm_s16le",
        output_path
    ])
    
    result = subprocess.run(cmd, capture_output=True)
    if result.returncode != 0:
        return None
    
    return output_path


def run_vad_classification(audio_path: str) -> tuple[str, float, int]:
    """
    Run Silero VAD to classify audio.
    
    Returns:
        (classification, speech_ratio, speech_segments_count)
        classification: one of "music_only", "speech_only", "speech_with_music"
    """
    if not HAS_TORCH:
        return ("speech_only", 0.5, 0)  # Default fallback
    
    try:
        # Read WAV file using scipy (avoids torchaudio dependency issues)
        sample_rate, wav_data = wavfile.read(audio_path)
        
        # Convert to float32 and normalize
        if wav_data.dtype == np.int16:
            wav = torch.from_numpy(wav_data.astype(np.float32) / 32768.0)
        elif wav_data.dtype == np.int32:
            wav = torch.from_numpy(wav_data.astype(np.float32) / 2147483648.0)
        elif wav_data.dtype == np.float32:
            wav = torch.from_numpy(wav_data)
        else:
            wav = torch.from_numpy(wav_data.astype(np.float32))
        
        # Ensure mono
        if len(wav.shape) > 1:
            wav = wav[:, 0]
        
        # Resample to 16kHz if needed
        if sample_rate != 16000:
            import torchaudio
            resampler = torchaudio.transforms.Resample(sample_rate, 16000)
            wav = resampler(wav)
            sample_rate = 16000
        
        # Load VAD model
        model, utils = torch.hub.load(
            repo_or_dir="snakers4/silero-vad",
            model="silero_vad",
            force_reload=False,
            trust_repo=True,
        )
        model.eval()
        
        get_speech_timestamps = utils[0]
        
        speech_timestamps = get_speech_timestamps(
            wav, model,
            sampling_rate=16000,
            min_speech_duration_ms=500,
            min_silence_duration_ms=300,
            return_seconds=True,
        )
        
        total_duration = len(wav) / 16000.0
        speech_duration = sum(ts["end"] - ts["start"] for ts in speech_timestamps)
        speech_ratio = speech_duration / total_duration if total_duration > 0 else 0.0
        
        # Classification logic:
        # - speech_ratio < 0.01: music_only (no speech detected)
        # - speech_ratio >= 0.01 and speech_ratio < 0.3: speech_with_music (sparse speech)
        # - speech_ratio >= 0.3: speech_only (mostly speech)
        
        if speech_ratio < 0.01:
            classification = "music_only"
        elif speech_ratio < 0.3:
            classification = "speech_with_music"
        else:
            classification = "speech_only"
        
        return (classification, speech_ratio, len(speech_timestamps))
        
    except Exception as e:
        print(f"VAD error: {e}", file=sys.stderr)
        return ("speech_only", 0.5, 0)


def probe_audio_track(file_path: str) -> dict:
    """
    Main function: probe audio track and classify.
    
    Returns:
        dict with: classification, speech_ratio, speech_segments, duration, 
                   track_index, track_language, all_tracks
    """
    # Get audio tracks
    tracks = get_audio_tracks(file_path)
    
    if not tracks:
        return {
            "classification": "no_audio",
            "speech_ratio": 0.0,
            "speech_segments": 0,
            "duration": 0.0,
            "track_index": None,
            "track_language": None,
            "all_tracks": [],
        }
    
    # Select best track
    best_index = select_best_track(tracks)
    best_track = next((t for t in tracks if t["index"] == best_index), tracks[0])
    
    # Extract audio for VAD
    audio_path = extract_audio_for_vad(file_path, best_index)
    
    if audio_path is None:
        return {
            "classification": "silent_audio",
            "speech_ratio": 0.0,
            "speech_segments": 0,
            "duration": 0.0,
            "track_index": best_index,
            "track_language": best_track.get("language", ""),
            "all_tracks": tracks,
        }
    
    # Get duration
    probe_cmd = [
        "ffprobe", "-v", "quiet",
        "-print_format", "json",
        "-show_format",
        audio_path
    ]
    probe_result = subprocess.run(probe_cmd, capture_output=True, text=True)
    duration = 0.0
    if probe_result.returncode == 0:
        probe_data = json.loads(probe_result.stdout)
        duration = float(probe_data.get("format", {}).get("duration", 0))
    
    # Run VAD
    classification, speech_ratio, speech_segments = run_vad_classification(audio_path)
    
    # Cleanup temp file
    Path(audio_path).unlink(missing_ok=True)
    
    return {
        "classification": classification,
        "speech_ratio": round(speech_ratio, 4),
        "speech_segments": speech_segments,
        "duration": round(duration, 2),
        "track_index": best_index,
        "track_language": best_track.get("language", ""),
        "all_tracks": tracks,
    }


def main():
    parser = argparse.ArgumentParser(description="Audio track probe with VAD classification")
    parser.add_argument("--file", "-f", required=True, help="Video file path")
    parser.add_argument("--json", "-j", action="store_true", help="Output as JSON")
    args = parser.parse_args()
    
    result = probe_audio_track(args.file)
    
    if args.json:
        print(json.dumps(result, indent=2))
    else:
        print(result["classification"])


if __name__ == "__main__":
    main()