import gradio as gr
import yt_dlp
import torch
import numpy as np
from faster_whisper import WhisperModel
import os
os.system("pip install -q gradio torch whisper faster-whisper yt-dlp numpy")



# Load model sekali aja
model = WhisperModel("small", device="cpu", compute_type="float32")

# Fungsi untuk transkripsi dari file
def transcribe_audio(file):
    segments, _ = model.transcribe(file)
    transcript = "\n".join(segment.text for segment in segments)

    # Simpan ke file TXT
    file_path = "transcription.txt"
    with open(file_path, "w", encoding="utf-8") as f:
        f.write(transcript)

    return f"**Transcription:**\n{transcript}"

# Fungsi untuk ambil audio dari YouTube
def get_audio_from_youtube(url):
    ydl_opts = {
        "format": "bestaudio/best",
        "postprocessors": [{"key": "FFmpegExtractAudio", "preferredcodec": "mp3", "preferredquality": "192"}],
        "outtmpl": "temp_audio.%(ext)s",
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True)
        return "temp_audio.mp3"

# Fungsi untuk transkripsi dari YouTube
def transcribe_youtube(url):
    audio_file = get_audio_from_youtube(url)
    return transcribe_audio(audio_file)

#memastikan file ada sebelum digunakan
if not os.path.exists('transcription.txt'):
    with open("transcription.txt", "w", encoding="utf-8") as f:
        f.write("")

# Gradio UI
with gr.Blocks() as app:
    gr.Markdown("# 🎤 YouTube & Audio Transcriber with Whisper AI")
    
    with gr.Tab("Upload File"):
        audio_input = gr.File(label="Upload Audio File")
        file_transcribe_button = gr.Button("Transcribe")
        file_output = gr.Textbox(label="Transcription")
        download_file = gr.File(label="Download Transcription", value="transcription.txt")

    with gr.Tab("YouTube Video"):
        youtube_url = gr.Textbox(label="YouTube URL")
        yt_transcribe_button = gr.Button("Transcribe")
        yt_output = gr.Textbox(label="Transcription")

    file_transcribe_button.click(transcribe_audio, inputs=audio_input, outputs=[file_output, download_file])
    yt_transcribe_button.click(transcribe_youtube, inputs=youtube_url, outputs=yt_output)

# Run Gradio app
app.launch(server_name="0.0.0.0", server_port=7880)