// 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 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 replStream = document.getElementById("repl-stream");
const transcriptEl = document.getElementById("transcript");
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;
    replStream.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 = "";
    replStream.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 — terminal-style: prompt line, then output lines underneath
    // indented under the prompt with a tier-tag suffix as a Lisp comment.
    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" : "");

        // Render the input across one or more lines (preserve newlines).
        const inputLines = entry.input.split(/\n/);
        for (let i = 0; i < inputLines.length; i++) {
            const line = document.createElement("div");
            line.className = "prompt-line";
            const sigil = document.createElement("span");
            sigil.className = "sigil";
            sigil.textContent = i === 0 ? "λ>" : "..";
            line.appendChild(sigil);
            line.appendChild(document.createTextNode(inputLines[i]));
            block.appendChild(line);
        }

        for (const r of (entry.results || [])) {
            const row = document.createElement("div");
            row.className = "tier-output";
            const result = document.createElement("span");
            result.className = "tier-result";
            const text = (r.output != null ? r.output : "");
            const errText = r.error ? "error: " + r.error : "";
            result.textContent = (text + (errText && text ? "\n" : "") + errText).replace(/\n$/, "");
            const meta = document.createElement("span");
            meta.className = "tier-meta";
            meta.textContent = `${TIER_LABEL[r.tier] || r.tier} · ${(r.elapsed | 0)}ms`;
            row.appendChild(result);
            row.appendChild(meta);
            block.appendChild(row);
        }
        transcriptEl.appendChild(block);
    }
    // Body owns the scroll now; jump to the bottom after layout settles.
    // requestAnimationFrame lets the just-mounted DOM contribute to
    // scrollHeight before we measure — otherwise on first load the page
    // sticks at the top because the transcript hasn't been laid out yet.
    requestAnimationFrame(() => {
        window.scrollTo({ top: document.documentElement.scrollHeight, behavior: "instant" });
    });
}

// ─── History navigation (readline-style up/down) ────────────────────
// Each tab owns its own history; the active tab's draft is preserved so
// that walking back into history doesn't eat what the user was typing.
const history = { idx: null, draft: "" };

function historyEntries() {
    // Transcript's `input` fields, deduplicated against the immediate
    // predecessor — pressing up shouldn't make you hit the same line
    // twice in a row when you just submitted it.
    const tab = activeTab();
    if (!tab) return [];
    const out = [];
    for (const e of tab.transcript) {
        if (out.length && out[out.length - 1] === e.input) continue;
        out.push(e.input);
    }
    return out;
}

function historyPrev() {
    const entries = historyEntries();
    if (entries.length === 0) return;
    if (history.idx === null) {
        history.draft = inputEl.value;
        history.idx = entries.length - 1;
    } else if (history.idx > 0) {
        history.idx--;
    }
    inputEl.value = entries[history.idx];
    inputEl.setSelectionRange(inputEl.value.length, inputEl.value.length);
}

function historyNext() {
    const entries = historyEntries();
    if (history.idx === null) return;
    if (history.idx >= entries.length - 1) {
        history.idx = null;
        inputEl.value = history.draft;
        history.draft = "";
    } else {
        history.idx++;
        inputEl.value = entries[history.idx];
    }
    inputEl.setSelectionRange(inputEl.value.length, inputEl.value.length);
}

function resetHistory() {
    history.idx = null;
    history.draft = "";
}

// ─── 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 = "";
    resetHistory();
    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();
            }
        });
    }
}

// ─── Heap pressure indicator ───────────────────────────────────────
// Polls each loaded tier in the active tab and prints a compact
// "py 12M · c 32M · asm 4M" string next to the tabbar buttons.
const heapEl = document.createElement("span");
heapEl.className = "heap-pressure";
heapEl.title = "tier worker memory — \"reboot tier\" reclaims on demand";

function formatBytes(n) {
    if (n < 1024) return `${n}B`;
    if (n < 1024 * 1024) return `${(n / 1024).toFixed(0)}K`;
    return `${(n / (1024 * 1024)).toFixed(1)}M`;
}

async function pollHeap() {
    const tab = activeTab();
    if (!tab) return;
    const tiers = ["python", "c", "asm"];
    const parts = [];
    for (const t of tiers) {
        const k = workerKey(tab.id, t);
        const w = state.workers[k];
        if (!w) continue;
        const runId = ++state.nextRunId;
        const stats = await new Promise((resolve) => {
            const handler = (e) => {
                if (e.data.kind !== "heap" || e.data.runId !== runId) return;
                w.removeEventListener("message", handler);
                resolve(e.data.stats);
            };
            w.addEventListener("message", handler);
            w.postMessage({ kind: "heap", runId, tier: t });
            // Timeout safety — eval-busy workers won't reply.
            setTimeout(() => { w.removeEventListener("message", handler); resolve(null); }, 200);
        });
        if (stats && stats.used != null) {
            parts.push(`${t.slice(0, 3)} ${formatBytes(stats.used)}`);
        }
    }
    heapEl.textContent = parts.length ? parts.join(" · ") : "";
}
setInterval(pollHeap, 2000);

// ─── 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);
// Insert heap pressure indicator into the tabbar after the spacer.
document.getElementById("tabbar").insertBefore(heapEl, resetTierBtn);

inputEl.addEventListener("keydown", (e) => {
    if (e.key === "Enter" && !e.shiftKey) {
        e.preventDefault();
        sendInput();
        return;
    }
    // Up/Down navigate per-tab history when the caret is in the only
    // (or first/last) row — otherwise they belong to the textarea's
    // natural multi-line navigation.
    if (e.key === "ArrowUp") {
        const before = inputEl.value.substring(0, inputEl.selectionStart);
        if (!before.includes("\n")) {
            e.preventDefault();
            historyPrev();
        }
        return;
    }
    if (e.key === "ArrowDown") {
        const after = inputEl.value.substring(inputEl.selectionStart);
        if (!after.includes("\n")) {
            e.preventDefault();
            historyNext();
        }
        return;
    }
});
inputEl.addEventListener("input", () => {
    // Any keystroke that isn't an arrow drops history navigation —
    // edits are now the user's own draft, not the historical entry.
    if (history.idx !== null && document.activeElement === inputEl) {
        // We can't reliably distinguish arrow-induced updates here, so
        // we only invalidate when the buffer has actually diverged from
        // the historical entry the cursor was on.
        const entries = historyEntries();
        if (entries[history.idx] !== inputEl.value) {
            history.draft = inputEl.value;
            history.idx = null;
        }
    }
});
sendBtn.addEventListener("click", sendInput);
tierSelectEl.addEventListener("change", () => {
    const tab = activeTab();
    if (tab) { tab.tier = tierSelectEl.value; saveSoon(); }
});