import gradio as gr
import yt_dlp
import os
import zipfile
import ffmpeg
import json
from datetime import timedelta

def download_video(url, output_dir="downloads"):
    os.makedirs(output_dir, exist_ok=True)
    ydl_opts = {
        'format': 'bestaudio/best',
        'outtmpl': f'{output_dir}/%(title)s.%(ext)s',
        'postprocessors': [{'key': 'FFmpegExtractAudio', 'preferredcodec': 'wav', 'preferredquality': '192'}]
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info_dict = ydl.extract_info(url, download=True)
        return ydl.prepare_filename(info_dict).replace(info_dict['ext'], 'wav')

def parse_srt(srt_content):
    captions = []
    for block in srt_content.strip().split("\n\n"):
        parts = block.split("\n")
        if len(parts) >= 3:
            index = parts[0]
            time_range = parts[1]
            text = " ".join(parts[2:])
            start_time, end_time = time_range.split(" --> ")
            captions.append({"start": start_time, "end": end_time, "text": text})
    return captions

def srt_to_seconds(time_str):
    h, m, s = time_str.split(':')
    s, ms = s.split(',')
    return int(h) * 3600 + int(m) * 60 + int(s) + int(ms) / 1000

def cut_audio(audio_path, captions, output_dir="tts_dataset"):
    os.makedirs(output_dir, exist_ok=True)
    metadata = []
    for i, cap in enumerate(captions):
        start_sec = srt_to_seconds(cap['start'])
        end_sec = srt_to_seconds(cap['end'])
        output_path = os.path.join(output_dir, f"clip_{i}.wav")
        (
            ffmpeg
            .input(audio_path, ss=start_sec, to=end_sec)
            .output(output_path, format='wav')
            .run(quiet=True, overwrite_output=True)
        )
        metadata.append({"file": output_path, "text": cap['text']})
    return metadata

def zip_dataset(output_dir="tts_dataset", zip_name="tts_data.zip"):
    zip_path = os.path.join(output_dir, zip_name)
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, _, files in os.walk(output_dir):
            for file in files:
                zipf.write(os.path.join(root, file), arcname=file)
    return zip_path

def process_data(video_url, audio_file, caption_file):
    if video_url:
        audio_path = download_video(video_url)
    elif audio_file:
        audio_path = audio_file.name
    else:
        return "No valid input provided."
    
    captions = parse_srt(caption_file.read().decode("utf-8"))
    metadata = cut_audio(audio_path, captions)
    with open("tts_dataset/metadata.json", "w") as f:
        json.dump(metadata, f, indent=4, ensure_ascii=False)
    
    zip_path = zip_dataset()
    return zip_path

demo = gr.Interface(
    fn=process_data,
    inputs=[
        gr.Textbox(label="YouTube Video URL (optional)"),
        gr.File(label="Upload Audio File (optional)"),
        gr.File(label="Upload Caption File (SRT/VTT)")
    ],
    outputs=gr.File(label="Download Processed ZIP"),
    title="TTS Dataset Cutter",
    description="Upload a video/audio file + caption file (SRT/VTT) to extract speech clips for TTS training. Supports YouTube download."
)

demo.launch()