#!/usr/bin/env python3
"""AI Hub 등 스튜디오 녹음 데이터를 전화 통화 품질로 전처리하는 스크립트.

깨끗한 오디오에 PSTN 시뮬레이션(밴드패스 + 다운샘플링 + G.711 companding)을 적용하여
실제 통화 녹음과 유사한 학습 데이터를 생성한다.

Usage:
    # 단일 파일
    python scripts/preprocess_phone_audio.py data/aihub_raw/sample.wav

    # 디렉토리 일괄 처리
    python scripts/preprocess_phone_audio.py data/aihub_raw/ -o data/aihub_phone/

    # companding 방식 지정 (기본: random)
    python scripts/preprocess_phone_audio.py data/aihub_raw/ --companding alaw

    # 원본도 함께 복사 (원본+전화 혼합 학습용)
    python scripts/preprocess_phone_audio.py data/aihub_raw/ -o data/training/ --keep-original
"""

from __future__ import annotations

import argparse
import logging
import shutil
import sys
from pathlib import Path

import librosa
import soundfile as sf

# 프로젝트 루트를 sys.path에 추가
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from src.common.phone_simulator import CompandingType, PhoneSimulator

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("preprocess_phone_audio")

SUPPORTED_EXTENSIONS = {".wav", ".mp3", ".m4a", ".ogg", ".flac"}


def find_audio_files(input_path: Path) -> list[Path]:
    """입력 경로에서 오디오 파일 목록 반환."""
    if input_path.is_file():
        if input_path.suffix.lower() in SUPPORTED_EXTENSIONS:
            return [input_path]
        logger.warning(f"지원하지 않는 파일 형식: {input_path.suffix}")
        return []

    files = []
    for ext in SUPPORTED_EXTENSIONS:
        files.extend(input_path.rglob(f"*{ext}"))
    return sorted(files)


def process_file(
    input_file: Path,
    output_dir: Path,
    simulator: PhoneSimulator,
    input_root: Path,
    keep_original: bool = False,
) -> bool:
    """단일 파일을 전화 품질로 변환."""
    try:
        # 원본 디렉토리 구조 유지
        relative = input_file.relative_to(input_root)
        output_file = output_dir / relative.with_suffix(".wav")
        output_file.parent.mkdir(parents=True, exist_ok=True)

        # 오디오 로드 (mono, 원본 SR 유지)
        audio, sr = librosa.load(str(input_file), sr=None, mono=True)

        # 전화 품질 시뮬레이션 적용
        processed, new_sr = simulator.process(audio, sr)

        # phone_ 접두사로 저장
        phone_output = output_file.with_name(f"phone_{output_file.name}")
        sf.write(str(phone_output), processed, new_sr, subtype="PCM_16")

        # 원본도 복사 (혼합 학습용)
        if keep_original:
            orig_output = output_file.with_name(f"orig_{output_file.name}")
            shutil.copy2(str(input_file), str(orig_output))

        return True

    except Exception as e:
        logger.error(f"처리 실패 [{input_file.name}]: {e}")
        return False


def main():
    parser = argparse.ArgumentParser(
        description="스튜디오 녹음 → 전화 통화 품질 전처리",
    )
    parser.add_argument(
        "input",
        type=Path,
        help="입력 오디오 파일 또는 디렉토리",
    )
    parser.add_argument(
        "-o", "--output",
        type=Path,
        default=None,
        help="출력 디렉토리 (기본: {input}_phone/)",
    )
    parser.add_argument(
        "--companding",
        type=str,
        choices=["alaw", "ulaw", "random"],
        default="random",
        help="G.711 companding 방식 (기본: random — 파일마다 랜덤 선택)",
    )
    parser.add_argument(
        "--keep-original",
        action="store_true",
        help="원본 파일도 출력 디렉토리에 복사 (원본+전화 혼합 학습용)",
    )
    args = parser.parse_args()

    # 입력 경로 확인
    input_path = args.input.resolve()
    if not input_path.exists():
        logger.error(f"입력 경로가 존재하지 않습니다: {input_path}")
        sys.exit(1)

    # 출력 디렉토리 결정
    if args.output:
        output_dir = args.output.resolve()
    else:
        if input_path.is_file():
            output_dir = input_path.parent / f"{input_path.stem}_phone"
        else:
            output_dir = input_path.parent / f"{input_path.name}_phone"
    output_dir.mkdir(parents=True, exist_ok=True)

    # 입력 루트 (상대 경로 계산용)
    input_root = input_path if input_path.is_dir() else input_path.parent

    # 오디오 파일 탐색
    audio_files = find_audio_files(input_path)
    if not audio_files:
        logger.error("처리할 오디오 파일이 없습니다.")
        sys.exit(1)

    logger.info(f"오디오 파일 {len(audio_files)}개 발견")
    logger.info(f"출력 디렉토리: {output_dir}")
    logger.info(f"Companding: {args.companding}")
    if args.keep_original:
        logger.info("원본 파일도 함께 복사합니다")

    # 시뮬레이터 생성
    companding = CompandingType(args.companding)
    simulator = PhoneSimulator(companding=companding)

    # 일괄 처리
    success = 0
    fail = 0
    for i, audio_file in enumerate(audio_files, 1):
        logger.info(f"[{i}/{len(audio_files)}] {audio_file.name}")
        if process_file(audio_file, output_dir, simulator, input_root, args.keep_original):
            success += 1
        else:
            fail += 1

    logger.info(f"완료: 성공 {success}, 실패 {fail}, 전체 {len(audio_files)}")


if __name__ == "__main__":
    main()