#!/usr/bin/env python3
"""
Script para extrair audio de referencia de um video do YouTube.
Uso: python extract_voice_ref.py --url "https://youtube.com/..." --start 60 --end 75 --output voice_ref.wav
"""

import argparse
import subprocess
import os
import tempfile


def download_youtube_segment(url: str, start_seconds: int, end_seconds: int, output_path: str):
    """Baixa um segmento de video do YouTube e extrai o audio."""

    duration = end_seconds - start_seconds

    print(f"Baixando audio do YouTube...")
    print(f"  URL: {url}")
    print(f"  Inicio: {start_seconds}s")
    print(f"  Duracao: {duration}s")

    # Criar arquivo temporario
    temp_file = tempfile.NamedTemporaryFile(suffix='.wav', delete=False)
    temp_file.close()

    # Baixar com yt-dlp
    cmd = [
        'yt-dlp',
        '-x',  # Extrair audio
        '--audio-format', 'wav',
        '--postprocessor-args', f'ffmpeg:-ss {start_seconds} -t {duration}',
        '-o', temp_file.name.replace('.wav', '.%(ext)s'),
        url
    ]

    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        print(f"Erro no yt-dlp: {result.stderr}")
        raise RuntimeError("Falha ao baixar video")

    # Converter para formato correto (22050 Hz, mono)
    print("Convertendo audio para formato correto...")
    cmd_convert = [
        'ffmpeg', '-y',
        '-i', temp_file.name,
        '-ar', '22050',  # Sample rate para StyleTTS2
        '-ac', '1',      # Mono
        output_path
    ]

    subprocess.run(cmd_convert, capture_output=True)

    # Limpar arquivo temporario
    if os.path.exists(temp_file.name):
        os.unlink(temp_file.name)

    print(f"Audio salvo em: {output_path}")
    return output_path


def extract_from_local_video(video_path: str, start_seconds: int, end_seconds: int, output_path: str):
    """Extrai audio de um video local."""

    duration = end_seconds - start_seconds

    print(f"Extraindo audio do video...")
    print(f"  Video: {video_path}")
    print(f"  Inicio: {start_seconds}s")
    print(f"  Duracao: {duration}s")

    cmd = [
        'ffmpeg', '-y',
        '-i', video_path,
        '-ss', str(start_seconds),
        '-t', str(duration),
        '-ar', '22050',
        '-ac', '1',
        output_path
    ]

    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        print(f"Erro: {result.stderr}")
        raise RuntimeError("Falha ao extrair audio")

    print(f"Audio salvo em: {output_path}")
    return output_path


def main():
    parser = argparse.ArgumentParser(description='Extrair audio de referencia para clonagem de voz')
    parser.add_argument('--url', '-u', help='URL do YouTube')
    parser.add_argument('--video', '-v', help='Caminho do video local')
    parser.add_argument('--start', '-s', type=int, default=0, help='Segundo inicial')
    parser.add_argument('--end', '-e', type=int, default=15, help='Segundo final')
    parser.add_argument('--output', '-o', default='voice_ref.wav', help='Arquivo de saida')

    args = parser.parse_args()

    if not args.url and not args.video:
        parser.error("Especifique --url ou --video")

    if args.url:
        download_youtube_segment(
            url=args.url,
            start_seconds=args.start,
            end_seconds=args.end,
            output_path=args.output
        )
    else:
        extract_from_local_video(
            video_path=args.video,
            start_seconds=args.start,
            end_seconds=args.end,
            output_path=args.output
        )


if __name__ == '__main__':
    main()