import os
import gradio as gr
import tempfile
import yt_dlp
from pydub import AudioSegment
from groq import Groq
from huggingface_hub import InferenceClient

# ✅ Environment Variables
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
HUGGINGFACE_API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")

if not GROQ_API_KEY or not HUGGINGFACE_API_TOKEN:
    raise EnvironmentError("Please set GROQ_API_KEY and HUGGINGFACE_API_TOKEN.")

groq_client = Groq(api_key=GROQ_API_KEY)
hf_client = InferenceClient(token=HUGGINGFACE_API_TOKEN)

# ✅ Download YouTube Audio
def download_youtube_audio(youtube_url):
    with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tmp_file:
        ydl_opts = {
            "format": "bestaudio/best",
            "outtmpl": tmp_file.name,
            "quiet": True,
            "postprocessors": [{
                "key": "FFmpegExtractAudio",
                "preferredcodec": "mp3",
                "preferredquality": "192",
            }],
        }
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([youtube_url])
        return tmp_file.name

# ✅ Split long audio into chunks (max 5 mins)
def split_audio(file_path, max_duration_ms=5*60*1000):
    audio = AudioSegment.from_file(file_path)
    chunks = []
    for i in range(0, len(audio), max_duration_ms):
        chunk = audio[i:i + max_duration_ms]
        temp_chunk = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
        chunk.export(temp_chunk.name, format="mp3")
        chunks.append(temp_chunk.name)
    return chunks

# ✅ Transcribe with Groq (chunk-wise)
def transcribe_audio(audio_path):
    try:
        chunks = split_audio(audio_path)
        transcript = ""
        for i, chunk in enumerate(chunks):
            with open(chunk, "rb") as f:
                response = groq_client.audio.transcriptions.create(
                    model="whisper-large-v3",
                    file=f
                )
            transcript += response.text + "\n"
        return transcript.strip()
    except Exception as e:
        return f"❌ Error during transcription: {e}"

# ✅ Summarize
def summarize_text(text, lang):
    try:
        if lang == "English":
            model = "facebook/bart-large-cnn"
            prompt = f"Summarize the following text in English:\n\n{text}"
        else:
            model = "facebook/mbart-large-50-many-to-many-mmt"
            prompt = f"مندرجہ ذیل عبارت کا جامع اردو خلاصہ تحریر کریں:\n\n{text}"

        output = hf_client.text_generation(
            model=model,
            prompt=prompt,
            max_new_tokens=250,
            temperature=0.7,
        )
        return output
    except Exception as e:
        return f"❌ Error during summarization: {e}"

# ✅ Main Pipeline
def process_input(youtube_url, audio_file, lang):
    if youtube_url:
        audio_path = download_youtube_audio(youtube_url)
    elif audio_file:
        audio_path = audio_file
    else:
        return "❌ Please upload an audio file or paste a YouTube link.", "", ""

    transcript = transcribe_audio(audio_path)
    if transcript.startswith("❌"):
        return transcript, "", ""

    summary = summarize_text(transcript, lang)
    return "✅ Transcription Completed!", transcript, summary

# ✅ Gradio Interface
with gr.Blocks(title="🎧 Urdu/English Audio Summarizer") as app:
    gr.Markdown("## 🎧 Urdu & English Audio Summarizer\nUpload audio or paste YouTube link below:")

    with gr.Row():
        youtube_link = gr.Textbox(label="📺 YouTube Link (optional)")
        language_choice = gr.Dropdown(["English", "Urdu"], value="English", label="🌐 Summary Language")

    audio_input = gr.Audio(type="filepath", label="🎙️ Upload Audio File (optional)")
    btn = gr.Button("🚀 Transcribe & Summarize")

    status = gr.Textbox(label="Status")
    transcript_box = gr.Textbox(label="📝 Transcription", lines=8)
    summary_box = gr.Textbox(label="🧩 Summary", lines=8)

    btn.click(process_input, [youtube_link, audio_input, language_choice], [status, transcript_box, summary_box])

app.launch()