#!/usr/bin/env python3
"""
G1: Audio Preprocessing — Chuẩn hóa audio cho STT
- Convert → 16kHz Mono
- Compand filter: đẩy giọng nhỏ = giọng to (chống mất chữ khi Đại đức hạ giọng)
- Volume boost 2x
- Tối ưu cho Google Chirp API

Usage:
    python3 audio_preprocess.py input.mp3 -o output_clean.wav
    python3 audio_preprocess.py input.wav -o output_clean.wav --dry-run
"""

import argparse
import subprocess
import sys
import os
from pathlib import Path


def check_ffmpeg():
    """Verify ffmpeg is available"""
    try:
        subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("❌ ffmpeg not found. Install: sudo apt install ffmpeg")
        return False


def get_audio_info(input_path: str) -> dict:
    """Extract audio metadata using ffprobe"""
    cmd = [
        "ffprobe", "-v", "quiet",
        "-show_entries", "stream=codec_name,sample_rate,channels,duration,bit_rate",
        "-of", "default=noprint_wrappers=1",
        input_path
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    info = {}
    for line in result.stdout.strip().split("\n"):
        if "=" in line:
            k, v = line.split("=", 1)
            info[k] = v
    return info


def preprocess_audio(input_path: str, output_path: str, dry_run: bool = False):
    """
    Chuẩn hóa audio với tham số tối ưu cho giọng Myanmar Phật giáo:
    
    1. -ac 1          → Mono (triệt tiêu phase noise)
    2. -ar 16000      → 16kHz (chuẩn tối ưu Google STT)
    3. volume=2.0     → Boost toàn bộ
    4. compand        → Dynamic range compression:
       - attacks=0    → Phản ứng tức thì với thay đổi âm lượng
       - points=...   → Map -80dB→-90dB (silence→silence)
                              -45dB→-15dB (whisper→boosted)
                              -27dB→-9dB  (quiet→clear)
                              0dB→-7dB    (loud→controlled)
    """
    info = get_audio_info(input_path)
    print(f"📁 Input: {input_path}")
    print(f"   Codec: {info.get('codec_name', '?')}")
    print(f"   Sample rate: {info.get('sample_rate', '?')} Hz")
    print(f"   Channels: {info.get('channels', '?')}")
    print(f"   Duration: {float(info.get('duration', 0)):.1f}s")

    # Compand filter — tối ưu cho giọng tụng đều đều + đoạn hạ giọng
    compand_filter = (
        "volume=2.0,"
        "compand=attacks=0:"
        "points=-80/-900|-45/-15|-27/-9|0/-7"
    )

    cmd = [
        "ffmpeg",
        "-i", input_path,
        "-ac", "1",              # Mono
        "-ar", "16000",           # 16kHz
        "-filter:a", compand_filter,
        "-c:a", "pcm_s16le",     # PCM 16-bit (chuẩn cho STT API)
        "-y",                     # Overwrite output
        output_path
    ]

    print(f"\n🔧 Filter: {compand_filter}")
    print(f"📁 Output: {output_path}")

    if dry_run:
        print("🏁 DRY RUN — skipping execution")
        print(f"   Would run: {' '.join(cmd)}")
        return

    print("\n⏳ Processing...")
    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        print(f"❌ ffmpeg error:\n{result.stderr}")
        sys.exit(1)

    # Verify output
    out_info = get_audio_info(output_path)
    out_size = os.path.getsize(output_path) / (1024 * 1024)
    print(f"\n✅ Done!")
    print(f"   Output: {out_info.get('sample_rate', '?')} Hz, "
          f"{out_info.get('channels', '?')} channel(s)")
    print(f"   Size: {out_size:.1f} MB")
    print(f"   Duration: {float(out_info.get('duration', 0)):.1f}s")


def main():
    parser = argparse.ArgumentParser(
        description="G1: Audio Preprocessing cho Myanmar STT Pipeline"
    )
    parser.add_argument("input", help="Input audio file (.mp3, .wav, .m4a...)")
    parser.add_argument("-o", "--output", default=None,
                        help="Output path (default: <input>_clean.wav)")
    parser.add_argument("--dry-run", action="store_true",
                        help="Show command without executing")
    args = parser.parse_args()

    if not check_ffmpeg():
        sys.exit(1)

    if args.output is None:
        input_stem = Path(args.input).stem
        args.output = f"{input_stem}_clean.wav"

    # Ensure output directory exists
    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)

    preprocess_audio(args.input, args.output, args.dry_run)


if __name__ == "__main__":
    main()