#!/usr/bin/env python3
"""
G2: Audio Chunking — Cắt audio thành block 5 phút
- Mỗi block 300s (tối ưu cho API Chirp)
- Overlap 0.5s giữa các block (tránh mất chữ ở ranh giới)
- Đánh số tuần tự, dễ retry riêng lẻ nếu fail

Usage:
    python3 audio_chunk.py input.wav -d 300 -o chunks/
    python3 audio_chunk.py input.wav -d 120 -o segments_2min/ --overlap 0.3
"""

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


def get_duration(input_path: str) -> float:
    """Get audio duration in seconds"""
    cmd = [
        "ffprobe", "-v", "quiet",
        "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1",
        input_path
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    return float(result.stdout.strip())


def chunk_audio(
    input_path: str,
    output_dir: str,
    chunk_duration: float = 300.0,
    overlap: float = 0.5
):
    """
    Split audio into chunks with overlap.
    
    Args:
        input_path: Path to audio file
        output_dir: Output directory for chunks
        chunk_duration: Duration per chunk in seconds (default: 300s = 5min)
        overlap: Overlap between chunks in seconds (default: 0.5s)
    """
    total_duration = get_duration(input_path)
    name = Path(input_path).stem
    
    # Calculate number of chunks
    effective_duration = chunk_duration - overlap
    num_chunks = math.ceil(total_duration / effective_duration)
    
    print(f"📁 Input: {input_path}")
    print(f"   Total duration: {total_duration:.1f}s ({total_duration/60:.1f} min)")
    print(f"   Chunk duration: {chunk_duration}s ({chunk_duration/60:.0f} min)")
    print(f"   Overlap: {overlap}s")
    print(f"   Number of chunks: {num_chunks}")
    
    os.makedirs(output_dir, exist_ok=True)
    
    for i in range(num_chunks):
        start_time = i * effective_duration
        duration = min(chunk_duration, total_duration - start_time)
        
        output_file = os.path.join(output_dir, f"{name}_chunk_{i:03d}.wav")
        
        cmd = [
            "ffmpeg",
            "-i", input_path,
            "-ss", str(start_time),
            "-t", str(duration),
            "-acodec", "copy",      # Copy without re-encoding (fast)
            "-y",
            output_file
        ]
        
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode != 0:
            print(f"❌ Chunk {i} failed:\n{result.stderr}")
            sys.exit(1)
        
        chunk_end = start_time + duration
        print(f"   ✅ Chunk {i:03d}: {start_time:.1f}s → {chunk_end:.1f}s "
              f"({duration:.1f}s) → {output_file}")

    print(f"\n✅ All {num_chunks} chunks created in {output_dir}/")


def main():
    parser = argparse.ArgumentParser(
        description="G2: Audio Chunking for Myanmar STT Pipeline"
    )
    parser.add_argument("input", help="Input audio file (.wav)")
    parser.add_argument("-d", "--duration", type=float, default=300.0,
                        help="Chunk duration in seconds (default: 300 = 5 min)")
    parser.add_argument("-o", "--output-dir", default="chunks/",
                        help="Output directory for chunks")
    parser.add_argument("--overlap", type=float, default=0.5,
                        help="Overlap between chunks in seconds (default: 0.5)")
    args = parser.parse_args()

    chunk_audio(args.input, args.output_dir, args.duration, args.overlap)


if __name__ == "__main__":
    main()