import json
import os
import shutil
import threading
import time
from pathlib import Path

import gradio as gr
import librosa
import torch
from huggingface_hub import snapshot_download
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor


MODEL_ID = os.getenv("MODEL_ID", "Kiragu/whisper-small-kikuyu-v5")
MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "96"))

_model = None
_processor = None
_model_lock = threading.Lock()


def normalize_words(value):
    cleaned = []

    for character in value.lower():
        if character.isalnum() or character in {"'", " "}:
            cleaned.append(character)
        else:
            cleaned.append(" ")

    return [word for word in "".join(cleaned).split() if word]


def levenshtein_distance(a, b):
    previous = list(range(len(b) + 1))

    for index_a, word_a in enumerate(a, start=1):
        diagonal = previous[0]
        previous[0] = index_a

        for index_b, word_b in enumerate(b, start=1):
            saved = previous[index_b]
            cost = 0 if word_a == word_b else 1
            previous[index_b] = min(
                previous[index_b] + 1,
                previous[index_b - 1] + 1,
                diagonal + cost,
            )
            diagonal = saved

    return previous[-1]


def compare_transcripts(reference, transcript):
    expected = normalize_words(reference)
    actual = normalize_words(transcript)

    if not expected:
        return "Paste a reference transcript to calculate match and WER."

    distance = levenshtein_distance(expected, actual)
    wer = round((distance / len(expected)) * 100)
    match = max(0, 100 - wer)

    return f"Match: {match}% | WER: {wer}% | Reference words: {len(expected)} | Model words: {len(actual)}"


def patch_tokenizer_config(snapshot_path):
    snapshot_path = Path(snapshot_path)
    config_path = snapshot_path / "tokenizer_config.json"

    if not config_path.exists():
        return snapshot_path

    data = json.loads(config_path.read_text(encoding="utf-8"))

    if not isinstance(data.get("extra_special_tokens"), list):
        return snapshot_path

    try:
        data.pop("extra_special_tokens")
        config_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
        return snapshot_path
    except OSError:
        patched_path = Path("/tmp/kikuyu-whisper-patched")

        if patched_path.exists():
            shutil.rmtree(patched_path)

        shutil.copytree(snapshot_path, patched_path, symlinks=True)
        patched_config = patched_path / "tokenizer_config.json"
        patched_data = json.loads(patched_config.read_text(encoding="utf-8"))
        patched_data.pop("extra_special_tokens", None)
        patched_config.write_text(json.dumps(patched_data, ensure_ascii=False, indent=2), encoding="utf-8")
        return patched_path


def load_model():
    global _model, _processor

    with _model_lock:
        if _model is not None and _processor is not None:
            return _model, _processor

        snapshot_path = snapshot_download(MODEL_ID)
        model_path = patch_tokenizer_config(snapshot_path)

        device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.float16 if device == "cuda" else torch.float32

        processor = AutoProcessor.from_pretrained(model_path)
        model = AutoModelForSpeechSeq2Seq.from_pretrained(model_path, torch_dtype=dtype)
        model.to(device)
        model.eval()

        _model = model
        _processor = processor
        return _model, _processor


def transcribe_audio(audio_path, reference_text):
    if not audio_path:
        return "", "Add or record an audio sample first.", "No audio provided."

    started = time.time()
    model, processor = load_model()
    device = next(model.parameters()).device

    audio_array, sampling_rate = librosa.load(audio_path, sr=16000, mono=True)
    inputs = processor(audio_array, sampling_rate=sampling_rate, return_tensors="pt")
    input_features = inputs.input_features.to(device)

    with torch.inference_mode():
        predicted_ids = model.generate(
            input_features,
            max_new_tokens=MAX_NEW_TOKENS,
            num_beams=1,
        )

    transcript = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0].strip()
    elapsed = round(time.time() - started, 1)
    metrics = compare_transcripts(reference_text or "", transcript)
    status = f"Done in {elapsed}s using {MODEL_ID}."

    return transcript, metrics, status


with gr.Blocks(title="Kikuyu Speech To Text") as demo:
    gr.Markdown("# Kikuyu Speech To Text")
    gr.Markdown("Upload or record Kikuyu audio, then compare the model transcript against a known Thiomi reference.")

    with gr.Row():
        audio_input = gr.Audio(
            label="Kikuyu audio",
            sources=["microphone", "upload"],
            type="filepath",
        )
        reference_input = gr.Textbox(
            label="Correct transcript",
            lines=7,
            placeholder="Optional: paste the Thiomi reference transcript here",
        )

    transcribe_button = gr.Button("Transcribe")

    transcript_output = gr.Textbox(label="Model transcript", lines=7)
    metrics_output = gr.Textbox(label="Accuracy check", lines=2)
    status_output = gr.Textbox(label="Status", lines=2)

    transcribe_button.click(
        fn=transcribe_audio,
        inputs=[audio_input, reference_input],
        outputs=[transcript_output, metrics_output, status_output],
        api_name="transcribe",
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()