// wasm/repl/repl.js
// Lumbda interactive REPL with multi-tab sessions + encrypted persistence.

import { openVault } from "./crypto.js";

const TIER_LABEL = {
    python: "python (pyodide)",
    c: "c (emcc)",
    asm: "asm (wat)",
};

// ─── DOM ────────────────────────────────────────────────────────────
const lockScreen = document.getElementById("lock-screen");
const passwordEl = document.getElementById("password");
const unlockBtn = document.getElementById("unlock");
const freshBtn = document.getElementById("fresh");
const lockNote = document.getElementById("lock-note");

const tabbar = document.getElementById("tabbar");
const tabsEl = document.getElementById("tabs");
const newTabBtn = document.getElementById("new-tab");
const resetTierBtn = document.getElementById("reset-tier");
const clearLogBtn = document.getElementById("clear-log");
const cancelBtn = document.getElementById("cancel");
const lockBtn = document.getElementById("lock");

const transcriptWrap = document.getElementById("transcript-wrap");
const transcriptEl = document.getElementById("transcript");
const promptBar = document.getElementById("prompt-bar");
const inputEl = document.getElementById("input");
const tierSelectEl = document.getElementById("tier-select");
const sendBtn = document.getElementById("send");

// ─── State ──────────────────────────────────────────────────────────
const state = {
    vault: null,        // { read, write, destroy, vaultId } or null (ephemeral)
    tabs: [],           // [{ id, name, tier, transcript: [{kind, text, results?}] }]
    activeTabId: null,
    workers: {},        // { [tabId+":"+tier]: Worker }
    pending: {},        // same key → { runId, reject }
    nextRunId: 1,
    nextTabId: 1,
    saveTimer: null,
};

// ─── Vault ──────────────────────────────────────────────────────────
async function saveSoon() {
    if (!state.vault) return;
    if (state.saveTimer) clearTimeout(state.saveTimer);
    state.saveTimer = setTimeout(async () => {
        try {
            await state.vault.write({
                tabs: state.tabs,
                activeTabId: state.activeTabId,
                savedAt: Date.now(),
            });
        } catch (e) {
            console.error("vault write failed:", e);
        }
    }, 400);
}

async function tryUnlock() {
    const pw = passwordEl.value;
    if (!pw) { lockNote.textContent = "password required (or click ephemeral)"; return; }
    lockNote.textContent = "";
    try {
        state.vault = await openVault(pw);
        const data = await state.vault.read();
        if (data && data.__decryptionFailed) {
            // Wrong password for an existing vault under this id.
            lockNote.textContent = "vault exists but password is wrong";
            state.vault = null;
            return;
        }
        if (data && Array.isArray(data.tabs) && data.tabs.length > 0) {
            state.tabs = data.tabs;
            state.activeTabId = data.activeTabId || data.tabs[0].id;
            state.nextTabId = Math.max(...state.tabs.map((t) => t.id)) + 1;
        } else {
            state.tabs = [];
            state.activeTabId = null;
            state.nextTabId = 1;
        }
        enterRepl();
    } catch (e) {
        lockNote.textContent = "unlock failed: " + e.message;
    }
}

function enterEphemeral() {
    state.vault = null;
    state.tabs = [];
    state.activeTabId = null;
    state.nextTabId = 1;
    enterRepl();
}

function enterRepl() {
    lockScreen.hidden = true;
    tabbar.hidden = false;
    transcriptWrap.hidden = false;
    promptBar.hidden = false;
    if (state.tabs.length === 0) newTab();
    else renderAll();
    inputEl.focus();
}

function relock() {
    // Terminate all workers, clear in-memory state, return to lock screen.
    for (const k of Object.keys(state.workers)) {
        try { state.workers[k].terminate(); } catch {}
    }
    state.workers = {};
    state.pending = {};
    state.vault = null;
    state.tabs = [];
    state.activeTabId = null;
    transcriptEl.innerHTML = "";
    tabsEl.innerHTML = "";
    tabbar.hidden = true;
    transcriptWrap.hidden = true;
    promptBar.hidden = true;
    lockScreen.hidden = false;
    passwordEl.value = "";
    passwordEl.focus();
}

// ─── Tabs ───────────────────────────────────────────────────────────
function newTab(name) {
    const id = state.nextTabId++;
    name = name || `session ${id}`;
    state.tabs.push({ id, name, tier: "c", transcript: [] });
    state.activeTabId = id;
    renderAll();
    saveSoon();
}

function closeTab(id) {
    const idx = state.tabs.findIndex((t) => t.id === id);
    if (idx < 0) return;
    // Terminate this tab's workers.
    for (const k of Object.keys(state.workers)) {
        if (k.startsWith(id + ":")) {
            try { state.workers[k].terminate(); } catch {}
            delete state.workers[k];
        }
    }
    state.tabs.splice(idx, 1);
    if (state.activeTabId === id) {
        state.activeTabId = state.tabs.length ? state.tabs[Math.max(0, idx - 1)].id : null;
    }
    if (state.tabs.length === 0) newTab();
    else renderAll();
    saveSoon();
}

function setActiveTab(id) {
    state.activeTabId = id;
    renderAll();
    saveSoon();
}

function activeTab() {
    return state.tabs.find((t) => t.id === state.activeTabId);
}

// ─── Workers ────────────────────────────────────────────────────────
function workerKey(tabId, tier) { return `${tabId}:${tier}`; }

function spawnWorker(tabId, tier) {
    const w = new Worker(new URL("./worker.mjs", import.meta.url), { type: "module" });
    state.workers[workerKey(tabId, tier)] = w;
    return w;
}

function ensureWorker(tabId, tier) {
    const k = workerKey(tabId, tier);
    if (!state.workers[k]) spawnWorker(tabId, tier);
    return state.workers[k];
}

function evalInTier(tabId, tier, src) {
    return new Promise((resolve, reject) => {
        const w = ensureWorker(tabId, tier);
        const runId = state.nextRunId++;
        const pendingKey = workerKey(tabId, tier);
        state.pending[pendingKey] = { runId, reject };
        const handler = (e) => {
            if (e.data.runId !== runId) return;
            if (e.data.kind === "done") {
                w.removeEventListener("message", handler);
                delete state.pending[pendingKey];
                resolve(e.data.output);
            } else if (e.data.kind === "error") {
                w.removeEventListener("message", handler);
                delete state.pending[pendingKey];
                reject(new Error(e.data.message));
            }
        };
        w.addEventListener("message", handler);
        w.postMessage({ kind: "eval", runId, tier, src });
    });
}

function rebootTier(tabId, tier) {
    const k = workerKey(tabId, tier);
    if (state.workers[k]) {
        try { state.workers[k].terminate(); } catch {}
        delete state.workers[k];
    }
    if (state.pending[k]) {
        try { state.pending[k].reject(new Error("rebooted")); } catch {}
        delete state.pending[k];
    }
}

function cancelAllPendingInActiveTab() {
    const tab = activeTab();
    if (!tab) return;
    for (const k of Object.keys(state.pending)) {
        if (k.startsWith(tab.id + ":")) {
            try { state.pending[k].reject(new Error("cancelled")); } catch {}
            try { state.workers[k].terminate(); } catch {}
            delete state.workers[k];
            delete state.pending[k];
        }
    }
}

// ─── Render ─────────────────────────────────────────────────────────
function renderAll() {
    // Tabs
    tabsEl.innerHTML = "";
    for (const t of state.tabs) {
        const el = document.createElement("div");
        el.className = "tab" + (t.id === state.activeTabId ? " active" : "");
        const name = document.createElement("span");
        name.textContent = t.name;
        name.addEventListener("click", () => setActiveTab(t.id));
        name.addEventListener("dblclick", (e) => {
            e.stopPropagation();
            const newName = prompt("rename tab:", t.name);
            if (newName) { t.name = newName.slice(0, 40); renderAll(); saveSoon(); }
        });
        const close = document.createElement("span");
        close.className = "tab-close";
        close.textContent = "×";
        close.title = "close tab";
        close.addEventListener("click", (e) => { e.stopPropagation(); closeTab(t.id); });
        el.appendChild(name);
        el.appendChild(close);
        tabsEl.appendChild(el);
    }

    // Transcript
    transcriptEl.innerHTML = "";
    const tab = activeTab();
    if (!tab) return;
    tierSelectEl.value = tab.tier;
    for (const entry of tab.transcript) {
        const block = document.createElement("div");
        block.className = "entry" + (entry.kind === "error" ? " error" : "");
        const prompt = document.createElement("div");
        prompt.className = "prompt-line";
        prompt.textContent = "λ> " + entry.input;
        block.appendChild(prompt);
        for (const r of (entry.results || [])) {
            const row = document.createElement("div");
            row.className = "tier-output";
            const lab = document.createElement("div");
            lab.className = "tier-label";
            lab.textContent = TIER_LABEL[r.tier] || r.tier;
            const tm = document.createElement("span");
            tm.className = "tier-time";
            tm.textContent = `  ${(r.elapsed | 0)}ms`;
            lab.appendChild(tm);
            const out = document.createElement("div");
            out.className = "tier-result";
            out.textContent = r.output || (r.error ? "error: " + r.error : "");
            row.appendChild(lab);
            row.appendChild(out);
            block.appendChild(row);
        }
        transcriptEl.appendChild(block);
    }
    transcriptWrap.scrollTop = transcriptWrap.scrollHeight;
}

// ─── Input handling ─────────────────────────────────────────────────
async function sendInput() {
    const src = inputEl.value.trim();
    if (!src) return;
    const tab = activeTab();
    if (!tab) return;
    tab.tier = tierSelectEl.value;
    const tiers = tab.tier === "all" ? ["python", "c", "asm"] : [tab.tier];
    const entry = { input: src, results: [], kind: "ok" };
    tab.transcript.push(entry);
    inputEl.value = "";
    sendBtn.disabled = true;
    cancelBtn.disabled = false;
    renderAll();
    saveSoon();

    const startTimes = {};
    for (const t of tiers) startTimes[t] = performance.now();

    const promises = tiers.map((t) =>
        evalInTier(tab.id, t, src)
            .then((output) => ({ tier: t, output, elapsed: performance.now() - startTimes[t] }))
            .catch((e) => ({ tier: t, error: e.message, elapsed: performance.now() - startTimes[t] })));
    // Append results as they arrive so user sees them in race order.
    let remaining = promises.length;
    for (const p of promises) {
        p.then((r) => {
            entry.results.push(r);
            if (r.error) entry.kind = "error";
            renderAll();
            saveSoon();
            remaining--;
            if (remaining === 0) {
                sendBtn.disabled = false;
                cancelBtn.disabled = true;
                inputEl.focus();
            }
        });
    }
}

// ─── Wire up ────────────────────────────────────────────────────────
unlockBtn.addEventListener("click", tryUnlock);
freshBtn.addEventListener("click", enterEphemeral);
passwordEl.addEventListener("keydown", (e) => {
    if (e.key === "Enter") tryUnlock();
});

newTabBtn.addEventListener("click", () => newTab());
resetTierBtn.addEventListener("click", () => {
    const tab = activeTab();
    if (tab) rebootTier(tab.id, tab.tier);
});
clearLogBtn.addEventListener("click", () => {
    const tab = activeTab();
    if (tab) { tab.transcript = []; renderAll(); saveSoon(); }
});
cancelBtn.addEventListener("click", cancelAllPendingInActiveTab);
lockBtn.addEventListener("click", relock);

inputEl.addEventListener("keydown", (e) => {
    if (e.key === "Enter" && !e.shiftKey) {
        e.preventDefault();
        sendInput();
    }
});
sendBtn.addEventListener("click", sendInput);
tierSelectEl.addEventListener("change", () => {
    const tab = activeTab();
    if (tab) { tab.tier = tierSelectEl.value; saveSoon(); }
});