import os
import csv
import random
import time
import difflib
from collections import defaultdict
import numpy as np
import librosa
import torch
import spaces
import gradio as gr
from transformers import AutoProcessor, AutoModelForCTC

# Configuration & Constants
MODEL_ID = "njand/wav2vec2-xls-r-latin"
HF_TOKEN = os.getenv("HF_TOKEN")
CSV_PATH = "passages.csv"

# Global Custom CSS for Clean UI & Typography
CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');

* {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif !important;
}

.transcription-box {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif !important;
    white-space: pre-wrap !important;
    word-break: break-word !important;
    font-size: 1.15em !important;
    line-height: 1.8 !important;
    padding: 16px !important;
    border: 1px solid #e5e7eb !important;
    border-radius: 8px !important;
    background: #ffffff !important;
    min-height: 80px !important;
}
"""

# Load PyTorch Model & Processor globally on CPU at startup
print("Loading PyTorch model and processor...")
processor = AutoProcessor.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = AutoModelForCTC.from_pretrained(MODEL_ID, token=HF_TOKEN)
model.eval()
print("PyTorch model loaded successfully!")


# External CSV Passage Loader
def load_passages(csv_file=CSV_PATH):
    """Loads macronized passages from a CSV file with fallback defaults."""
    passages = []
    if os.path.exists(csv_file):
        try:
            with open(csv_file, mode="r", encoding="utf-8") as f:
                reader = csv.DictReader(f)
                for row in reader:
                    if "text" in row and "author" in row and "work" in row:
                        passages.append(row)
        except Exception as e:
            print(f"Error reading CSV: {e}")

    if not passages:
        passages = [
            {"author": "Caesar", "work": "De Bello Gallico", "text": "Gallia est omnis dīvīsa in partēs trēs."},
            {"author": "Cicero", "work": "In Catilinam", "text": "Quō usque tandem abutēre Catilīna patientiā nostrā?"},
            {"author": "Vergil", "work": "Aeneid", "text": "Arma virumque canō Trōiae quī prīmus ab ōrīs."}
        ]
    return passages

PASSAGES_LIST = load_passages()


# CTC Prefix Beam Search Implementation (Top-5)
def ctc_prefix_beam_search(logits: np.ndarray, beam_width: int = 10, top_n: int = 5):
    """
    Performs CTC Prefix Beam Search on model output logits.
    Returns top-N hypotheses with log probabilities.
    """
    if logits.ndim == 3:
        logits = logits[0]

    max_logits = np.max(logits, axis=-1, keepdims=True)
    exp_logits = np.exp(logits - max_logits)
    probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
    log_probs = np.log(probs + 1e-10)

    T, _ = log_probs.shape
    pad_id = getattr(processor.tokenizer, "pad_token_id", 0)

    beams = {(): (0.0, -float("inf"))}

    for t in range(T):
        next_beams = defaultdict(lambda: (-float("inf"), -float("inf")))
        top_indices = np.argsort(log_probs[t])[-beam_width:]

        for prefix, (p_b, p_nb) in beams.items():
            p_total = np.logaddexp(p_b, p_nb)

            for c in top_indices:
                pr = log_probs[t, c]

                if c == pad_id:
                    n_pb, n_pnb = next_beams[prefix]
                    next_beams[prefix] = (np.logaddexp(n_pb, p_total + pr), n_pnb)
                else:
                    last_char = prefix[-1] if len(prefix) > 0 else None
                    if c == last_char:
                        n_pb, n_pnb = next_beams[prefix + (c,)]
                        next_beams[prefix + (c,)] = (n_pb, np.logaddexp(n_pnb, p_b + pr))
                        
                        n_pb2, n_pnb2 = next_beams[prefix]
                        next_beams[prefix] = (n_pb2, np.logaddexp(n_pnb2, p_nb + pr))
                    else:
                        n_pb, n_pnb = next_beams[prefix + (c,)]
                        next_beams[prefix + (c,)] = (n_pb, np.logaddexp(n_pnb, p_total + pr))

        sorted_beams = sorted(
            next_beams.items(),
            key=lambda x: np.logaddexp(x[1][0], x[1][1]),
            reverse=True
        )[:beam_width]
        beams = dict(sorted_beams)

    decoded_results = []
    seen = set()

    for prefix, (p_b, p_nb) in beams.items():
        score = float(np.logaddexp(p_b, p_nb))
        text = processor.decode(list(prefix), skip_special_tokens=True).strip()
        if text and text not in seen:
            seen.add(text)
            decoded_results.append((text, score))
            if len(decoded_results) == top_n:
                break

    return decoded_results


# Primary HTML Transcription Renderer with Word-Space Support & Tooltips
def render_primary_transcription_html(logits: np.ndarray, reference_text: str = "", use_reference: bool = False):
    """
    Renders model transcription in clean typography with accurate word spaces.
    Highlights mismatches in soft red ONLY if reference alignment is enabled.
    Hovering always reveals confidence scores regardless of alignment settings.
    """
    if logits.ndim == 3:
        logits = logits[0]

    exp_logits = np.exp(logits - np.max(logits, axis=-1, keepdims=True))
    probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)

    pad_id = getattr(processor.tokenizer, "pad_token_id", 0)
    greedy_ids = np.argmax(probs, axis=-1)
    
    char_prob_pairs = []
    prev_id = None

    for t, token_id in enumerate(greedy_ids):
        if token_id != pad_id:
            raw_token = processor.tokenizer.convert_ids_to_tokens(int(token_id))
            if raw_token in ["|", " ", "<space>"]:
                char = " "
            else:
                char = processor.decode([token_id], skip_special_tokens=True)

            prob = float(probs[t, token_id])
            if char:
                if token_id == prev_id and char_prob_pairs:
                    last_c, last_p = char_prob_pairs[-1]
                    char_prob_pairs[-1] = (last_c, min(last_p, prob))
                else:
                    char_prob_pairs.append((char, prob))
        prev_id = token_id

    if not char_prob_pairs:
        return '<div class="transcription-box" style="color: #6b7280;">[No speech detected]</div>', "[No speech detected]"

    hyp_text = "".join([c for c, _ in char_prob_pairs])
    html_spans = []

    # Alignment matching if reference passage scoring is enabled
    if use_reference and reference_text.strip():
        ref_clean = reference_text.strip().lower()
        hyp_clean = hyp_text.lower()
        matcher = difflib.SequenceMatcher(None, hyp_clean, ref_clean)
        
        char_ref_map = {}
        for tag, i1, i2, j1, j2 in matcher.get_opcodes():
            if tag == "equal":
                for idx in range(i1, i2):
                    ref_idx = j1 + (idx - i1)
                    char_ref_map[idx] = ref_clean[ref_idx] if ref_idx < len(ref_clean) else None
            elif tag == "replace":
                for idx in range(i1, i2):
                    ref_idx = j1 + (idx - i1)
                    char_ref_map[idx] = ref_clean[ref_idx] if ref_idx < len(ref_clean) else "Mismatch"

        for idx, (c, p) in enumerate(char_prob_pairs):
            display_c = "&nbsp;" if c == " " else c
            expected_c = char_ref_map.get(idx)
            
            # Highlight STRICTLY on character mismatches (no low confidence highlighting)
            is_mismatch = (expected_c is not None and expected_c != c.lower())

            tooltip_parts = [f"Confidence: {p*100:.1f}%"]
            if expected_c:
                tooltip_parts.append(f"Expected: '{expected_c}'")
            tooltip_str = " | ".join(tooltip_parts)

            if is_mismatch:
                span_html = (
                    f'<span style="background-color: #fee2e2; color: #991b1b; border-bottom: 2px solid #f87171; '
                    f'padding: 0 1px; border-radius: 2px; font-weight: 600; cursor: help;" '
                    f'title="{tooltip_str}">{display_c}</span>'
                )
            else:
                span_html = f'<span style="color: #111827; cursor: help;" title="{tooltip_str}">{display_c}</span>'
            html_spans.append(span_html)

    else:
        # Standard view (No reference script): Zero highlighting, hovering reveals confidence
        for c, p in char_prob_pairs:
            display_c = "&nbsp;" if c == " " else c
            tooltip_str = f"Confidence: {p*100:.1f}%"
            span_html = f'<span style="color: #111827; cursor: help;" title="{tooltip_str}">{display_c}</span>'
            html_spans.append(span_html)

    rendered_html = f'<div class="transcription-box">{"".join(html_spans)}</div>'
    return rendered_html, hyp_text


# Audio Preprocessing
def preprocess_audio(audio_tuple):
    """Converts multi-channel PCM input to 16 kHz mono float32 array."""
    sample_rate, audio_data = audio_tuple

    if audio_data.ndim > 1:
        audio_data = audio_data.mean(axis=1)

    if np.issubdtype(audio_data.dtype, np.integer):
        max_val = np.iinfo(audio_data.dtype).max
        audio_data = audio_data.astype(np.float32) / max_val
    elif audio_data.dtype != np.float32:
        audio_data = audio_data.astype(np.float32)

    duration = len(audio_data) / sample_rate

    if sample_rate != 16000:
        audio_data = librosa.resample(audio_data, orig_sr=sample_rate, target_sr=16000)

    return audio_data, duration


# Main Inference Pipeline (ZeroGPU Accelerated)
@spaces.GPU
def transcribe_audio(audio_tuple, reference_text="", use_reference=False, progress=gr.Progress(track_tqdm=False)):
    """Runs preprocessing, PyTorch CUDA inference, CTC Beam Search, and HTML rendering."""
    if audio_tuple is None:
        return (
            '<div class="transcription-box" style="color: #ef4444;">⚠️ Please record or upload an audio file first.</div>',
            [],
            "No audio provided."
        )

    try:
        progress(0.15, desc="Preprocessing audio...")
        start_time = time.perf_counter()
        audio_data, duration = preprocess_audio(audio_tuple)

        progress(0.35, desc="Extracting features and attaching GPU...")
        model.to("cuda")
        inputs = processor(audio_data, sampling_rate=16000, return_tensors="pt").input_values.to("cuda")

        progress(0.60, desc="Running PyTorch model on GPU...")
        with torch.no_grad():
            logits = model(inputs).logits

        progress(0.80, desc="Executing CTC Prefix Beam Search (Top-5)...")
        logits_np = logits.detach().cpu().numpy()[0]
        beam_results = ctc_prefix_beam_search(logits_np, beam_width=10, top_n=5)

        end_time = time.perf_counter()
        inference_time = end_time - start_time
        rtf = inference_time / duration if duration > 0 else 0.0

        if not beam_results:
            beam_table_data = []
        else:
            beam_table_data = [[i + 1, res[0], f"{res[1]:.2f}"] for i, res in enumerate(beam_results)]

        metrics_md = (
            f"- **Audio Duration:** {duration:.2f} seconds\n"
            f"- **Inference Time:** {inference_time:.3f} seconds\n"
            f"- **Real-Time Factor (RTF):** {rtf:.3f} "
            f"({'⚡ Faster than real-time' if rtf < 1.0 else '🐌 Slower than real-time'})"
        )

        progress(0.95, desc="Rendering transcription alignment...")
        primary_html, _ = render_primary_transcription_html(
            logits_np, reference_text=reference_text, use_reference=use_reference
        )

        progress(1.0, desc="Done!")
        return primary_html, beam_table_data, metrics_md

    except Exception as e:
        err_msg = f'<div class="transcription-box" style="color: #ef4444;">❌ An error occurred: {str(e)}</div>'
        return err_msg, [], f"Error: {str(e)}"


# Gradio UI Construction
with gr.Blocks(title="Latin Speech Recognition & Pronunciation Demo", css=CUSTOM_CSS) as demo:
    
    # State Management
    playlist_indices = list(range(len(PASSAGES_LIST)))
    random.shuffle(playlist_indices)
    playlist_state = gr.State(playlist_indices)
    pointer_state = gr.State(0)

    # Header Section
    gr.Markdown(
        f"""
        # 🎙️ Latin Automatic Speech Recognition (ASR)
        
        **Model:** [`{MODEL_ID}`](https://huggingface.co/{MODEL_ID})
        """
    )

    with gr.Row():
        # Left Panel: Input & Alignment Controls
        with gr.Column(scale=1):
            audio_input = gr.Audio(
                sources=["microphone", "upload"],
                type="numpy",
                label="Audio Input",
            )

            # Step 2 Grouped Single Container
            with gr.Group():
                use_passage_cb = gr.Checkbox(
                    label="Reading Passage for Pronunciation Scoring",
                    value=False
                )

                # Dependent Reference Passage Elements (Hidden by Default)
                with gr.Group(visible=False) as passage_container:
                    with gr.Row():
                        prev_btn = gr.Button("◀ Previous", size="sm", variant="secondary")
                        next_btn = gr.Button("Next ▶", size="sm", variant="secondary")

                    passage_display = gr.Markdown(value="*Click 'Next Passage' to load a text.*")
                    
                    reference_input = gr.Textbox(
                        label="Reference Text",
                        placeholder="Type custom text or click 'Next Passage'...",
                        lines=2,
                        interactive=True
                    )

            with gr.Row():
                clear_btn = gr.Button("Clear All", variant="secondary")
                submit_btn = gr.Button("Transcribe Audio", variant="primary")

        # Right Panel: Results & Metrics
        with gr.Column(scale=1):
            with gr.Tabs():
                with gr.Tab("📝 Primary Result"):
                    gr.Markdown("**Transcription View:** Hover over letters to inspect model confidence scores & alignment.")
                    primary_output = gr.HTML(
                        value='<div class="transcription-box" style="color: #6b7280;">Your transcription will appear here...</div>'
                    )

                with gr.Tab("📊 Top-5 Beam Search"):
                    beam_table = gr.Dataframe(
                        headers=["Rank", "Hypothesis Text", "Log Score"],
                        label="Top-5 Candidates",
                        interactive=False,
                    )

                with gr.Tab("⚡ Inference Metrics"):
                    metrics_output = gr.Markdown(
                        value="Metrics will update after audio is processed."
                    )

    # Passage Navigation & Interactive Visibility Logic
    def toggle_passage_visibility(is_enabled):
        return gr.update(visible=is_enabled)

    use_passage_cb.change(
        fn=toggle_passage_visibility,
        inputs=use_passage_cb,
        outputs=passage_container
    )

    def get_passage_at_index(playlist, pointer):
        if not playlist:
            return "*No passages loaded.*", ""
        idx = playlist[pointer % len(playlist)]
        item = PASSAGES_LIST[idx]
        display_md = f"**Author:** {item['author']} — *{item['work']}*"
        return display_md, item["text"]

    def navigate_next(playlist, pointer):
        new_pointer = pointer + 1
        disp, txt = get_passage_at_index(playlist, new_pointer)
        return new_pointer, disp, txt

    def navigate_prev(playlist, pointer):
        new_pointer = pointer - 1
        disp, txt = get_passage_at_index(playlist, new_pointer)
        return new_pointer, disp, txt

    def handle_custom_text_input(text):
        """Clears classical author citation when user types custom text."""
        if not text.strip():
            return "*Click 'Next Passage' to load a text.*"
        for item in PASSAGES_LIST:
            if item["text"].strip() == text.strip():
                return f"**Author:** {item['author']} — *{item['work']}*"
        return "*(Custom Reference Text)*"

    reference_input.input(
        fn=handle_custom_text_input,
        inputs=reference_input,
        outputs=passage_display
    )

    next_btn.click(
        fn=navigate_next,
        inputs=[playlist_state, pointer_state],
        outputs=[pointer_state, passage_display, reference_input]
    )

    prev_btn.click(
        fn=navigate_prev,
        inputs=[playlist_state, pointer_state],
        outputs=[pointer_state, passage_display, reference_input]
    )

    submit_btn.click(
        fn=transcribe_audio,
        inputs=[audio_input, reference_input, use_passage_cb],
        outputs=[primary_output, beam_table, metrics_output],
    )

    clear_btn.click(
        fn=lambda: (
            None, False, gr.update(visible=False), "", "*Click 'Next Passage' to load a text.*",
            '<div class="transcription-box" style="color: #6b7280;">Your transcription will appear here...</div>',
            [], "Metrics will update after audio is processed."
        ),
        inputs=None,
        outputs=[audio_input, use_passage_cb, passage_container, reference_input, passage_display, primary_output, beam_table, metrics_output],
    )

demo.launch(theme=gr.themes.Soft())