import os
import shutil
import tempfile
from pathlib import Path

import gradio as gr
from transformers import pipeline

# Prøv å støtte video via pydub + ffmpeg hvis tilgjengelig (valgfritt)
try:
    from pydub import AudioSegment
    HAS_PYDUB = True
except Exception:
    HAS_PYDUB = False

# --- Konfigurasjon ---
# CPU: bruk en mindre, flerspråklig modell. (large-v3 på CPU vil ofte knekke.)
# Alternativer: "openai/whisper-small", "openai/whisper-medium", "distil-whisper/distil-small.multilingual"
ASR_MODEL_ID = "openai/whisper-small"

def make_transcriber():
    # device=-1 tvinger CPU. return_timestamps=True gir tidskoder i retur.
    return pipeline(
        task="automatic-speech-recognition",
        model=ASR_MODEL_ID,
        device=-1,
        return_timestamps=True
    )

transcriber = make_transcriber()

VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".webm"}
AUDIO_EXTS = {".wav", ".mp3", ".m4a", ".flac", ".ogg", ".opus", ".aac"}

def extract_audio_if_needed(input_path: str) -> str:
    """
    Tar inn en filsti (audio eller video).
    Hvis video og pydub+ffmpeg finnes, ekstraheres WAV til temp-katalog og
    vi returnerer ny filsti. Hvis ikke, kastes en forklarende feil.
    Hvis allerede audio, returneres originalstien.
    """
    suffix = Path(input_path).suffix.lower()

    # Allerede audio?
    if suffix in AUDIO_EXTS:
        return input_path

    # Video?
    if suffix in VIDEO_EXTS:
        if not HAS_PYDUB:
            raise RuntimeError(
                "Video oppdaget, men pydub/ffmpeg er ikke tilgjengelig. "
                "Installer pydub og ffmpeg (se requirements.txt og apt.txt), "
                "eller last opp en ren lydfil."
            )
        # Ekstraher WAV
        temp_dir = tempfile.mkdtemp(prefix="asr_")
        out_wav = os.path.join(temp_dir, "extracted_audio.wav")
        audio = AudioSegment.from_file(input_path)
        audio.export(out_wav, format="wav")
        return out_wav

    # Ukjent – la Whisper prøve; hvis det feiler, får brukeren feilmelding
    return input_path

def handle_upload_and_transcribe(file_path: str):
    if not file_path:
        return "Last opp en lyd- eller videofil."

    tmp_to_cleanup = None
    try:
        # Kan generere en temp WAV (for video)
        maybe_audio = extract_audio_if_needed(file_path)
        if maybe_audio != file_path:
            tmp_to_cleanup = os.path.dirname(maybe_audio)

        # Tips: du kan sette språk eksplisitt for raskere/mer stabil dekoding:
        # generate_kwargs={"task": "transcribe", "language": "no"}
        result = transcriber(maybe_audio)

        # Rydd temp
        if tmp_to_cleanup and os.path.exists(tmp_to_cleanup):
            shutil.rmtree(tmp_to_cleanup, ignore_errors=True)

        # Normaliser utdata
        if isinstance(result, dict):
            # transformers>=4.30 gir ofte {"text": "...", "chunks": [...]}
            text = result.get("text")
            if text:
                return text.strip()
            # fallback
            return str(result)
        elif isinstance(result, list) and result:
            return result[0].get("text", str(result))
        return str(result)

    except Exception as e:
        # Rydd opp ved feil
        if tmp_to_cleanup and os.path.exists(tmp_to_cleanup):
            shutil.rmtree(tmp_to_cleanup, ignore_errors=True)
        return f"❌ Feil under prosessering/transkripsjon: {e}"

with gr.Blocks(title="Multilingual Audio/Video Transcription") as demo:
    gr.Markdown(
        "## Multilingual Transcription (CPU)\n"
        "Last opp en lydfil (.wav/.mp3/.m4a/…) eller videofil (.mp4/.mov/…). "
        "På CPU brukes en mindre Whisper-modell for stabil kjøring."
    )
    inp = gr.Audio(type="filepath", label="Fil (audio eller video)")
    out = gr.Textbox(label="Transkripsjon")
    btn = gr.Button("Transkriber")
    btn.click(handle_upload_and_transcribe, inputs=inp, outputs=out)

if __name__ == "__main__":
    # På HF Spaces trenger du vanligvis ikke server_name/server_port her.
    demo.launch()