import gradio as gr
from pydub import AudioSegment
from google import genai
import tempfile
import os
import time
import re
import concurrent.futures

API_TIMEOUT_SEC = 300

PROMPT = """
Transcribe the following audio exactly as spoken. 
Format the transcript with clear paragraphs. 
When the speaker moves to a new topic or subject, insert a subheading in markdown (### Topic) that summarizes the upcoming content. 
For verses, stanzas, or ślokas, break them into separate lines as they would appear in a written poem (one line per quarter or line as appropriate). 
Maintain the original language rules:
- Sanskrit proper names: IAST without italics.
- Other Sanskrit words and phrases: IAST with *italics* and English translation in parentheses, e.g., *śāntiḥ* (*peace*).
- English: exactly as spoken.
- Serbian/Croatian: word order as given (e.g., 'Svāmī je klimnuo glavom' not 'klimnuo je').
Do not add timestamps. Output only the formatted transcript.
"""

def call_gemini(seg_path, api_key):
    """Poziva Gemini API (izvršava se u pozadinskoj niti)."""
    client = genai.Client(api_key=api_key)
    myfile = client.files.upload(file=seg_path)
    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=[myfile, PROMPT]
    )
    text = response.text
    client.files.delete(name=myfile.name)
    return text

def transcribe_segment(audio_path, api_key, start_min, end_min, progress=gr.Progress()):
    # Provera ulaza
    if audio_path is None:
        yield "❌ Niste otpremili MP3 fajl!"
        return
    if not api_key.strip():
        yield "❌ Unesite Gemini API ključ."
        return
    if start_min is None or end_min is None:
        yield "❌ Unesite početak i kraj u minutima."
        return
    if start_min >= end_min:
        yield "❌ Početak mora biti manji od kraja."
        return

    progress(0.0, desc="Učitavanje audio fajla...")
    yield "⏳ Učitavanje audio fajla..."
    try:
        sound = AudioSegment.from_file(audio_path)
    except Exception as e:
        yield f"❌ Greška pri učitavanju fajla: {e}"
        return

    total_sec = sound.duration_seconds
    start_ms = int(start_min * 60 * 1000)
    end_ms = int(end_min * 60 * 1000)

    if end_ms > len(sound):
        yield f"❌ Kraj ({end_min} min) premašuje trajanje snimka ({total_sec/60:.1f} min)."
        return
    if start_ms < 0:
        start_ms = 0

    segment = sound[start_ms:end_ms]
    segment_dur = len(segment) / 1000
    progress(0.1, desc=f"Isečeno {start_min}-{end_min} min ({segment_dur/60:.1f} min)")
    yield f"✂️ Isečen segment {start_min}–{end_min} min ({segment_dur/60:.1f} min)"

    progress(0.2, desc="Priprema audio segmenta...")
    yield "📦 Priprema audio segmenta..."
    tmp = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
    segment.export(tmp.name, format="mp3")
    seg_path = tmp.name

    api_key_clean = api_key.strip()
    last_err = ""
    for attempt in range(1, 10):
        progress(0.25 + attempt * 0.07, desc=f"Transkripcija (pokušaj {attempt})...")
        yield f"🎤 Transkripcija u toku (pokušaj {attempt})..."
        print(f"API poziv za {start_min}-{end_min} min, pokušaj {attempt}")
        try:
            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
                future = executor.submit(call_gemini, seg_path, api_key_clean)
                text = future.result(timeout=API_TIMEOUT_SEC)
            progress(1.0, desc="Gotovo!")
            os.unlink(seg_path)
            yield f"### {start_min}–{end_min} min\n{text}"
            return
        except concurrent.futures.TimeoutError:
            last_err = "API poziv je prekoračio vremensko ograničenje (5 min)."
            print("   Timeout!")
            break
        except Exception as e:
            last_err = str(e)
            print(f"   Greška: {last_err}")
            if "503" in last_err or "UNAVAILABLE" in last_err:
                wait = 10 * attempt
                print(f"   503, čekam {wait}s...")
                yield f"⏳ Server nedostupan (503), čekam {wait}s..."
                time.sleep(wait)
            elif "429" in last_err and "RESOURCE_EXHAUSTED" in last_err:
                match = re.search(r"retryDelay':\s*'?(\d+)s?'?", last_err)
                wait = int(match.group(1)) + 2 if match else 60
                print(f"   429, čekam {wait}s...")
                yield f"⏳ Previše zahteva (429), čekam {wait}s..."
                time.sleep(wait)
            else:
                break

    # Sve probe neuspešne
    try:
        os.unlink(seg_path)
    except:
        pass
    yield f"### {start_min}–{end_min} min\n*Greška: {last_err}*"

with gr.Blocks() as demo:
    gr.Markdown("# 🎙️ Transkripcija po opsezima (progres uživo + formatiranje)")

    with gr.Row():
        audio_in = gr.Audio(type="filepath", label="Otpremite MP3 fajl")
        api_key = gr.Textbox(label="Gemini API ključ", type="password")

    with gr.Row():
        start_min = gr.Number(label="Početak (minuti)", value=0, precision=1)
        end_min = gr.Number(label="Kraj (minuti)", value=20, precision=1)

    gr.Markdown("*Otpremite ceo MP3, pa unosite opsege jedan po jedan (npr. 0–20, 20–40…).*")

    btn = gr.Button("Pokreni transkripciju ovog opsega")
    output = gr.Markdown(label="Transkript")

    btn.click(
        fn=transcribe_segment,
        inputs=[audio_in, api_key, start_min, end_min],
        outputs=output
    )

if __name__ == "__main__":
    demo.launch()