# app.py
import time
import traceback
from datetime import datetime

import torch
import gradio as gr
import requests
from transformers import pipeline
from transformers.pipelines.audio_utils import ffmpeg_read

# ----------------------------
# Firebase (PUBLIC DB)
# ----------------------------
FIREBASE_URL = "https://speechad-32698-default-rtdb.firebaseio.com"

def firebase_put_transcribe(text: str) -> bool:
    """
    Overwrite /transcribe with latest transcription
    """
    payload = {
        "text": text,
        "updated_at": datetime.utcnow().isoformat() + "Z",
        "length": len(text),
    }
    try:
        url = f"{FIREBASE_URL}/transcribe.json"
        r = requests.put(url, json=payload, timeout=10)
        print("Firebase PUT status:", r.status_code)
        print("Firebase PUT response:", r.text)
        return r.ok
    except Exception as e:
        print("Firebase PUT exception:", e)
        return False

# ----------------------------
# Whisper model
# ----------------------------
MODEL_NAME = "openai/whisper-small"
device = 0 if torch.cuda.is_available() else -1

pipe = pipeline(
    "automatic-speech-recognition",
    model=MODEL_NAME,
    device=device,
    chunk_length_s=60,
    ignore_warning=True,
)

CHUNK_SEC = 2
OVERLAP_SEC = 0.05
SLEEP_BETWEEN_CHUNKS = 0.01

def merge_overlap(prev: str, new: str, max_words: int = 12) -> str:
    if not prev:
        return new
    pw = prev.split()
    nw = new.split()
    max_k = min(max_words, len(pw), len(nw))
    for k in range(max_k, 0, -1):
        if pw[-k:] == nw[:k]:
            return " ".join(pw + nw[k:])
    return prev + " " + new

# ----------------------------
# Transcription
# ----------------------------
def transcribe_file(filepath: str):
    if not filepath:
        raise gr.Error("No audio provided")

    try:
        sr = pipe.feature_extractor.sampling_rate

        with open(filepath, "rb") as f:
            audio = ffmpeg_read(f.read(), sr)

        accumulated = ""
        chunk_samples = int(CHUNK_SEC * sr)
        overlap_samples = int(OVERLAP_SEC * sr)

        for start in range(0, len(audio), chunk_samples):
            end = min(len(audio), start + chunk_samples)
            chunk = audio[max(0, start - overlap_samples):min(len(audio), end + overlap_samples)]

            result = pipe({"array": chunk, "sampling_rate": sr})
            text = result.get("text", "").strip()
            accumulated = merge_overlap(accumulated, text)

            yield accumulated, gr.update(visible=False)
            time.sleep(SLEEP_BETWEEN_CHUNKS)

        # FINAL WRITE (overwrite)
        ok = firebase_put_transcribe(accumulated)

        if ok:
            yield accumulated, gr.update(
                visible=True,
                value="""
                <div style="text-align:center;margin-top:20px">
                  <a href="https://nolist-testingspeechwhisper.hf.space"
                     style="padding:12px 24px;
                            background:#2563eb;
                            color:white;
                            border-radius:8px;
                            text-decoration:none;
                            font-weight:600;">
                    Continue
                  </a>
                </div>
                """
            )
        else:
            yield accumulated, gr.update(
                visible=True,
                value="<b style='color:red'>Firebase write failed</b>"
            )

    except Exception as e:
        raise gr.Error(f"Transcription failed: {e}\n{traceback.format_exc()}")

# ----------------------------
# Gradio UI
# ----------------------------
with gr.Blocks(title="Record Option") as demo:
    gr.Markdown("# Record Option")

    audio = gr.Audio(sources=["microphone"], type="filepath")
    output = gr.Textbox(lines=12, label="Transcription")
    html = gr.HTML(visible=False)

    audio.change(transcribe_file, audio, [output, html])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)