import gradio as gr
import yt_dlp
import whisper
import os

# 1. Load Whisper Model (Small model CPU par fast chalta hai)
# Agar GPU available hai, to ye automatically use karega, warna CPU.
print("Loading Whisper Model...")
model = whisper.load_model("base") 
print("Model Loaded!")

def get_audio_from_tiktok(url):
    """
    TikTok URL se audio download karne ka function using yt-dlp
    """
    try:
        # Output filename template
        output_filename = "downloaded_audio"
        
        # Agar purani file hai to delete karein
        if os.path.exists(f"{output_filename}.mp3"):
            os.remove(f"{output_filename}.mp3")

        ydl_opts = {
            'format': 'bestaudio/best',
            'outtmpl': output_filename,  # File name without extension
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
                'preferredquality': '192',
            }],
            'quiet': True,
            'no_warnings': True,
            'user_agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }

        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([url])
        
        return f"{output_filename}.mp3"
    
    except Exception as e:
        return str(e)

def process_tiktok(tiktok_url):
    """
    Main function jo UI se connect hoga
    """
    if not tiktok_url:
        return "Error: Please enter a valid URL."

    # Step 1: Download Audio
    print(f"Downloading from: {tiktok_url}")
    audio_path = get_audio_from_tiktok(tiktok_url)
    
    # Check if download was successful (audio path should be a file path, not error text)
    if not audio_path.endswith(".mp3"):
        return f"Download Failed: {audio_path}"

    # Step 2: Transcribe using Whisper
    print("Transcribing...")
    try:
        # Whisper audio ko text me badal dega
        result = model.transcribe(audio_path)
        transcript = result["text"]
        return transcript
    except Exception as e:
        return f"Transcription Error: {str(e)}"

# --- Gradio UI ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎵 TikTok to Text Transcriber
        Paste a TikTok link below to get the text transcript of the video.
        """
    )
    
    with gr.Row():
        inp_url = gr.Textbox(label="TikTok Video URL", placeholder="Paste link here (e.g., https://www.tiktok.com/@user/video/...)")
        btn = gr.Button("Transcribe 📝", variant="primary")
    
    out_text = gr.Textbox(label="Transcript", lines=10, show_copy_button=True)

    btn.click(fn=process_tiktok, inputs=inp_url, outputs=out_text)

# Launch
demo.launch()