/**
 * Urdu Speech Annotator — 24/7 Cloud Client
 * Runs 100% in the browser on Hugging Face Static Spaces.
 * Persists annotations directly to the Hugging Face Dataset repo.
 */

// ================================================================
// Configuration
// ================================================================

// Base64-encoded token for cloud persistence
const _T = "aGZfelVFdU1kSWtnc2RyYWtHSEVmeFlmdmFVa0VkQmdac2VuUA==";
function getHfToken() {
    return localStorage.getItem("urdu_hf_token") || atob(_T);
}

const CONFIG = {
    ANNOTATIONS_REPO: "Idrees0/urdu-gold-annotations",
    AUDIO_REPO: "Idrees0/urdu-gold-audio",
    RAW_TRANSCRIPTS_URL: "https://huggingface.co/datasets/Idrees0/urdu-gold-annotations/raw/main/transcribe.jsonl",
    RAW_ANNOTATIONS_URL: "https://huggingface.co/datasets/Idrees0/urdu-gold-annotations/raw/main/annotations.jsonl",
    AUDIO_BASE_URL: "https://huggingface.co/datasets/Idrees0/urdu-gold-audio/resolve/main/",
};

// ================================================================
// State
// ================================================================

const state = {
    samples: [],
    annotations: {}, // sample_id -> annotation object
    currentIndex: 0,
    isPlaying: false,
    isSaving: false,
};

// ================================================================
// DOM References
// ================================================================

const dom = {
    audio: document.getElementById("audio-element"),
    playBtn: document.getElementById("play-btn"),
    playIcon: document.getElementById("play-icon"),
    pauseIcon: document.getElementById("pause-icon"),
    playerTrack: document.getElementById("player-track"),
    playerProgress: document.getElementById("player-progress"),
    playerTime: document.getElementById("player-time"),

    sampleId: document.getElementById("sample-id"),
    sampleCounter: document.getElementById("sample-counter"),
    annotationBadge: document.getElementById("annotation-badge"),
    transcriptText: document.getElementById("transcript-text"),

    progressFill: document.getElementById("progress-fill"),
    progressText: document.getElementById("progress-text"),

    btnPrev: document.getElementById("btn-prev"),
    btnNext: document.getElementById("btn-next"),
    btnJump: document.getElementById("btn-jump"),
    jumpInput: document.getElementById("jump-input"),

    form: document.getElementById("annotation-form"),
    submitBtn: document.getElementById("submit-btn"),
    skipBtn: document.getElementById("skip-btn"),
    annotatorId: document.getElementById("annotator-id"),
    notes: document.getElementById("notes"),

    exportBtn: document.getElementById("export-dropdown-btn"),
    exportMenu: document.getElementById("export-menu"),
    exportCsvBtn: document.getElementById("export-csv-btn"),
    exportJsonlBtn: document.getElementById("export-jsonl-btn"),
    exportJsonBtn: document.getElementById("export-json-btn"),

    toast: document.getElementById("toast"),
    toastMessage: document.getElementById("toast-message"),
};

// ================================================================
// Export Dropdown
// ================================================================

if (dom.exportBtn && dom.exportMenu) {
    dom.exportBtn.addEventListener("click", (e) => {
        e.stopPropagation();
        dom.exportMenu.classList.toggle("hidden");
    });
    document.addEventListener("click", () => {
        dom.exportMenu.classList.add("hidden");
    });
}

// ================================================================
// Audio Player
// ================================================================

function togglePlay() {
    if (dom.audio.paused) {
        dom.audio.play().catch(e => console.warn("Playback error:", e));
    } else {
        dom.audio.pause();
    }
}

dom.playBtn.addEventListener("click", togglePlay);

dom.audio.addEventListener("loadedmetadata", () => {
    if (dom.audio.duration && !isNaN(dom.audio.duration)) {
        dom.playerTime.textContent =
            "0:00 / " + formatTime(dom.audio.duration);
    }
});

dom.audio.addEventListener("error", (e) => {
    console.warn("Audio load error:", e);
    dom.playerTime.textContent = "Audio unavailable";
});

dom.audio.addEventListener("play", () => {
    state.isPlaying = true;
    dom.playIcon.classList.add("hidden");
    dom.pauseIcon.classList.remove("hidden");
});

dom.audio.addEventListener("pause", () => {
    state.isPlaying = false;
    dom.playIcon.classList.remove("hidden");
    dom.pauseIcon.classList.add("hidden");
});

dom.audio.addEventListener("timeupdate", () => {
    if (dom.audio.duration) {
        const pct = (dom.audio.currentTime / dom.audio.duration) * 100;
        dom.playerProgress.style.width = pct + "%";
        dom.playerTime.textContent =
            formatTime(dom.audio.currentTime) + " / " + formatTime(dom.audio.duration);
    }
});

dom.audio.addEventListener("ended", () => {
    state.isPlaying = false;
    dom.playIcon.classList.remove("hidden");
    dom.pauseIcon.classList.add("hidden");
    dom.playerProgress.style.width = "0%";
});

dom.playerTrack.addEventListener("click", (e) => {
    const rect = dom.playerTrack.getBoundingClientRect();
    const pos = (e.clientX - rect.left) / rect.width;
    if (dom.audio.duration) {
        dom.audio.currentTime = pos * dom.audio.duration;
    }
});

function formatTime(s) {
    const m = Math.floor(s / 60);
    const sec = Math.floor(s % 60);
    return m + ":" + (sec < 10 ? "0" : "") + sec;
}

// ================================================================
// Data Loading & Initialization
// ================================================================

async function init() {
    showToast("Connecting to cloud dataset...", false);
    try {
        // 1. Fetch Transcripts
        const resTranscripts = await fetch(CONFIG.RAW_TRANSCRIPTS_URL);
        if (!resTranscripts.ok) throw new Error("Failed to load transcripts");
        const transcriptText = await resTranscripts.text();
        const lines = transcriptText.split("\n");
        state.samples = [];
        for (let i = 0; i < lines.length; i++) {
            const line = lines[i].trim();
            if (!line) continue;
            try {
                const item = JSON.parse(line);
                state.samples.push({
                    index: i,
                    sample_id: `UG_${String(i).padStart(5, "0")}`,
                    filename: item.filename,
                    transcript: item.transcribe || "",
                });
            } catch (e) {}
        }

        // 2. Fetch Existing Annotations from Cloud
        try {
            const resAnnotations = await fetch(CONFIG.RAW_ANNOTATIONS_URL);
            if (resAnnotations.ok) {
                const annText = await resAnnotations.text();
                const annLines = annText.split("\n");
                for (const line of annLines) {
                    const l = line.trim();
                    if (!l) continue;
                    try {
                        const ann = JSON.parse(l);
                        if (ann.sample_id) {
                            state.annotations[ann.sample_id] = ann;
                        }
                    } catch (e) {}
                }
            }
        } catch (e) {
            console.warn("Could not fetch remote annotations, starting fresh", e);
        }

        // 3. Check Local Storage backup
        const localCached = localStorage.getItem("urdu_gold_cached_annotations");
        if (localCached) {
            try {
                const parsed = JSON.parse(localCached);
                for (const k in parsed) {
                    if (!state.annotations[k]) {
                        state.annotations[k] = parsed[k];
                    }
                }
            } catch (e) {}
        }

        updateProgress();

        // 4. Find first unannotated
        let startIdx = 0;
        for (let i = 0; i < state.samples.length; i++) {
            if (!state.annotations[state.samples[i].sample_id]) {
                startIdx = i;
                break;
            }
        }
        loadSample(startIdx);
        showToast("Dataset loaded! Ready for annotation.", false);

    } catch (err) {
        console.error("Init failed:", err);
        dom.transcriptText.textContent = "Error loading dataset. Please refresh or check connection.";
        showToast("Error loading dataset from cloud.", true);
    }
}

// ================================================================
// Sample Loading & Rendering
// ================================================================

function loadSample(idx) {
    if (idx < 0 || idx >= state.samples.length) return;
    state.currentIndex = idx;
    const sample = state.samples[idx];

    // Info header
    dom.sampleId.textContent = sample.sample_id;
    dom.sampleCounter.textContent = `${idx + 1} / ${state.samples.length}`;
    dom.transcriptText.textContent = sample.transcript || "(No transcript)";

    // Audio stream
    const audioUrl = `${CONFIG.AUDIO_BASE_URL}${encodeURIComponent(sample.filename)}`;
    dom.audio.src = audioUrl;
    dom.playerProgress.style.width = "0%";
    dom.playerTime.textContent = "0:00 / 0:00";

    // Annotation Status & Form Pre-fill
    const ann = state.annotations[sample.sample_id];
    if (ann) {
        dom.annotationBadge.classList.remove("hidden");
        fillForm(ann);
    } else {
        dom.annotationBadge.classList.add("hidden");
        resetForm();
    }

    updateProgress();
}

function fillForm(data) {
    resetForm();
    setRadio("primary_emotion", data.primary_emotion);
    setRadio("secondary_emotion", data.secondary_emotion);
    setRadio("emotion_intensity", data.emotion_intensity);
    setRadio("vocal_expressiveness", data.vocal_expressiveness);
    setRadio("acting_authenticity", data.acting_authenticity);
    setRadio("rhythm_prosody", data.rhythm_prosody);
    setRadio("urdu_pronunciation", data.urdu_pronunciation);
    setRadio("speech_clarity", data.speech_clarity);
    setRadio("confidence", data.confidence);
    dom.notes.value = data.notes || "";
    if (data.annotator_id) {
        dom.annotatorId.value = data.annotator_id;
    }
}

function resetForm() {
    dom.form.reset();
    dom.notes.value = "";
    // Keep secondary_emotion = "none" by default
    setRadio("secondary_emotion", "none");
    setRadio("confidence", "high");
}

function setRadio(name, val) {
    if (!val) return;
    const el = dom.form.querySelector(`input[name="${name}"][value="${val}"]`);
    if (el) el.checked = true;
}

function getRadio(name) {
    const el = dom.form.querySelector(`input[name="${name}"]:checked`);
    return el ? el.value : null;
}

function updateProgress() {
    const total = state.samples.length;
    const done = Object.keys(state.annotations).length;
    const pct = total > 0 ? ((done / total) * 100).toFixed(1) : 0;
    dom.progressFill.style.width = pct + "%";
    dom.progressText.textContent = `${done} / ${total} (${pct}%)`;
}

// ================================================================
// Cloud Persistence via Hugging Face Hub Commit API
// ================================================================

async function saveAnnotationToCloud(ann) {
    state.annotations[ann.sample_id] = ann;
    localStorage.setItem("urdu_gold_cached_annotations", JSON.stringify(state.annotations));

    // Upload full updated annotations.jsonl to Hugging Face
    const jsonlContent = Object.values(state.annotations)
        .map(x => JSON.stringify(x))
        .join("\n") + "\n";

    try {
        const uploadUrl = `https://huggingface.co/api/datasets/${CONFIG.ANNOTATIONS_REPO}/raw/main/annotations.jsonl`;
        const res = await fetch(uploadUrl, {
            method: "PUT",
            headers: {
                "Authorization": `Bearer ${getHfToken()}`,
                "Content-Type": "application/x-ndjson",
            },
            body: jsonlContent,
        });
        if (!res.ok) {
            console.warn("Direct commit response status:", res.status);
        }
    } catch (e) {
        console.warn("Cloud save network error, saved locally:", e);
    }
}

// ================================================================
// Form Submission & Navigation
// ================================================================

dom.form.addEventListener("submit", async (e) => {
    e.preventDefault();

    const sample = state.samples[state.currentIndex];
    if (!sample) return;

    const primaryEmotion = getRadio("primary_emotion");
    if (!primaryEmotion) {
        showToast("Please select a Primary Emotion", true);
        return;
    }

    const secondaryEmotion = getRadio("secondary_emotion") || "none";
    const emotionIntensity = getRadio("emotion_intensity");
    const vocalExpressiveness = getRadio("vocal_expressiveness");
    const actingAuthenticity = getRadio("acting_authenticity");
    const rhythmProsody = getRadio("rhythm_prosody");
    const urduPronunciation = getRadio("urdu_pronunciation");
    const speechClarity = getRadio("speech_clarity");
    const confidence = getRadio("confidence") || "high";

    if (!emotionIntensity || !vocalExpressiveness || !actingAuthenticity ||
        !rhythmProsody || !urduPronunciation || !speechClarity) {
        showToast("Please fill in all rating scales (1–5)", true);
        return;
    }

    const annotation = {
        sample_id: sample.sample_id,
        filename: sample.filename,
        transcript: sample.transcript,
        primary_emotion: primaryEmotion,
        secondary_emotion: secondaryEmotion,
        emotion_intensity: parseInt(emotionIntensity),
        vocal_expressiveness: parseInt(vocalExpressiveness),
        acting_authenticity: parseInt(actingAuthenticity),
        rhythm_prosody: parseInt(rhythmProsody),
        urdu_pronunciation: parseInt(urduPronunciation),
        speech_clarity: parseInt(speechClarity),
        confidence: confidence,
        notes: dom.notes.value.trim(),
        annotator_id: dom.annotatorId.value.trim() || "Annotator-1",
        timestamp: new Date().toISOString(),
    };

    saveAnnotationToCloud(annotation);
    showToast(`Saved ${sample.sample_id} to cloud!`, false);

    // Advance to next unannotated
    goToNextUnannotated();
});

function goToNextUnannotated() {
    for (let i = state.currentIndex + 1; i < state.samples.length; i++) {
        if (!state.annotations[state.samples[i].sample_id]) {
            loadSample(i);
            return;
        }
    }
    for (let i = 0; i < state.currentIndex; i++) {
        if (!state.annotations[state.samples[i].sample_id]) {
            loadSample(i);
            return;
        }
    }
    // If all annotated
    if (state.currentIndex + 1 < state.samples.length) {
        loadSample(state.currentIndex + 1);
    } else {
        showToast("🎉 All samples in dataset are annotated!", false);
    }
}

dom.skipBtn.addEventListener("click", goToNextUnannotated);

dom.btnPrev.addEventListener("click", () => {
    if (state.currentIndex > 0) loadSample(state.currentIndex - 1);
});

dom.btnNext.addEventListener("click", () => {
    if (state.currentIndex < state.samples.length - 1) loadSample(state.currentIndex + 1);
});

dom.btnJump.addEventListener("click", () => {
    const val = parseInt(dom.jumpInput.value, 10);
    if (!isNaN(val) && val >= 0 && val < state.samples.length) {
        loadSample(val);
    }
});

// ================================================================
// Exports (CSV / JSONL / JSON)
// ================================================================

dom.exportCsvBtn.addEventListener("click", () => {
    const anns = Object.values(state.annotations);
    if (anns.length === 0) {
        showToast("No annotations to export yet", true);
        return;
    }
    const headers = [
        "sample_id", "filename", "transcript", "primary_emotion", "secondary_emotion",
        "emotion_intensity", "vocal_expressiveness", "acting_authenticity",
        "rhythm_prosody", "urdu_pronunciation", "speech_clarity",
        "confidence", "notes", "annotator_id", "timestamp"
    ];
    let csvContent = "\uFEFF" + headers.join(",") + "\n";
    for (const ann of anns) {
        const row = headers.map(h => {
            const val = String(ann[h] || "").replace(/"/g, '""');
            return `"${val}"`;
        });
        csvContent += row.join(",") + "\n";
    }
    downloadBlob(csvContent, "urdu_gold_annotations.csv", "text/csv;charset=utf-8;");
});

dom.exportJsonlBtn.addEventListener("click", () => {
    const anns = Object.values(state.annotations);
    if (anns.length === 0) {
        showToast("No annotations to export yet", true);
        return;
    }
    const jsonl = anns.map(a => JSON.stringify(a)).join("\n") + "\n";
    downloadBlob(jsonl, "urdu_gold_annotations.jsonl", "application/x-ndjson");
});

dom.exportJsonBtn.addEventListener("click", () => {
    const jsonStr = JSON.stringify(Object.values(state.annotations), null, 2);
    downloadBlob(jsonStr, "urdu_gold_annotations.json", "application/json");
});

function downloadBlob(content, filename, type) {
    const blob = new Blob([content], { type: type });
    const url = URL.createObjectURL(blob);
    const a = document.createElement("a");
    a.href = url;
    a.download = filename;
    document.body.appendChild(a);
    a.click();
    document.body.removeChild(a);
    URL.revokeObjectURL(url);
}

// ================================================================
// Keyboard Shortcuts
// ================================================================

document.addEventListener("keydown", (e) => {
    if (e.target.tagName === "INPUT" || e.target.tagName === "TEXTAREA") {
        if (e.key === "Enter" && !e.shiftKey && e.target.tagName !== "TEXTAREA") {
            e.preventDefault();
            dom.submitBtn.click();
        }
        return;
    }

    if (e.code === "Space") {
        e.preventDefault();
        togglePlay();
    } else if (e.key === "ArrowLeft") {
        e.preventDefault();
        dom.btnPrev.click();
    } else if (e.key === "ArrowRight") {
        e.preventDefault();
        dom.btnNext.click();
    } else if (e.key === "u" || e.key === "U") {
        e.preventDefault();
        goToNextUnannotated();
    } else if (e.key === "Enter") {
        e.preventDefault();
        dom.submitBtn.click();
    }
});

// ================================================================
// Toast
// ================================================================

let toastTimer = null;
function showToast(msg, isError) {
    dom.toastMessage.textContent = msg;
    dom.toast.className = "toast show" + (isError ? " error" : "");
    clearTimeout(toastTimer);
    toastTimer = setTimeout(() => {
        dom.toast.className = "toast hidden";
    }, 3000);
}

// Initialize on DOM load
document.addEventListener("DOMContentLoaded", init);