// ==UserScript==
// @name         OpenRouter Chat Enhancements
// @namespace    http://tampermonkey.net/
// @license      MIT
// @version      1.2.0
// @description  Navigation hotkeys, message highlight, floating speaker, scroll protections, perfect collapse/expand handling, and enhanced edit scroll lock.
// @author       Rekt
// @match        https://openrouter.ai/chat*
// @grant        GM_addStyle
// @grant        GM_setValue
// @grant        GM_getValue
// @grant        GM_registerMenuCommand
// @run-at       document-end
// @downloadURL none
// ==/UserScript==

(function () {
    'use strict';

    /*********************** CONFIG & STYLES **********************/
    GM_addStyle(`
        /* Message highlight */
        .openrouter-nav-highlight {
            outline: 2.5px solid #22caff !important;
            z-index: 10 !important;
            border-radius: 7px;
            box-shadow: 0 0 0 3px #22caff33, 0 0 4px #22caff44;
            transition: outline 0.17s;
        }
        /* Floating Speaker */
        #openrouter-nav-speaker-float {
            position: fixed;
            top: 8px;
            left: 50%;
            transform: translateX(-50%);
            z-index: 999999;
            background: rgba(20,24,42,0.97);
            color: #f1f1f1;
            font-weight: 600;
            font-size: 1.05rem;
            padding: 7px 24px 7px 14px;
            border-radius: 19px;
            opacity: 0;
            visibility: hidden;
            pointer-events: none;
            min-width: 88px;
            display: flex;
            align-items: center;
            gap: 11px;
            box-shadow: 0 6px 17px #0002;
            max-width: 89vw;
            line-height: 1.2;
            transition: opacity 0.21s, top 0.15s;
        }
        #openrouter-nav-speaker-float.openrouter-nav-visible {
            opacity: 1;
            visibility: visible;
        }
        #openrouter-nav-speaker-float img {
            width: 27px;
            height: 27px;
            border-radius: 50%;
            border: 1px solid #fff6;
            object-fit: cover;
            margin-right: 6px;
        }
        #openrouter-nav-speaker-float span {
            white-space: nowrap;
            overflow: hidden;
            text-overflow: ellipsis;
            max-width: 320px;
        }
        /* Floating control panel */
        #openrouter-nav-panel {
            position: fixed;
            right: 20px;
            bottom: 25px;
            z-index: 999998;
            background: rgba(30,33,40,0.92);
            border-radius: 14px;
            box-shadow: 0 4px 22px #0003;
            padding: 11px 12px 11px 11px;
            display: flex;
            gap: 7px;
            align-items: center;
            transition: opacity 0.22s;
            user-select: none;
        }
        .openrouter-nav-btn {
            width: 42px;
            height: 42px;
            display: flex;
            align-items: center;
            justify-content: center;
            border: none;
            outline: none;
            background: rgba(245,245,255,0.08);
            border-radius: 6px;
            color: #fff;
            font-size: 18px;
            transition: background .14s;
            cursor: pointer;
        }
        .openrouter-nav-btn svg {
            width: 21px;
            height: 21px;
        }
        .openrouter-nav-btn:hover {
            background: rgba(245,245,255,0.19);
        }
        .openrouter-nav-divider {
            width: 7px;
        }
        @media (max-width: 768px) {
            #openrouter-nav-panel {
                right: 6px;
                bottom: 9px;
                padding: 7.5px 7px 7.5px 6px;
            }
            .openrouter-nav-btn {
                width: 34px;
                height: 34px;
            }
        }
    `);

    // Settings with persistence
    let modifierKey = GM_getValue('or_modifierKey', "Alt"); // "Alt", "Ctrl", "None"
    let panelEnabled = GM_getValue('or_panelEnabled', true);

    // Menu commands
    GM_registerMenuCommand("Set Hotkey Modifier: (Alt/Ctrl/None)", () => {
        const val = prompt('Use which key as the hotkey modifier? (Alt, Ctrl, None)', modifierKey);
        if (!val) return;
        const normalized = val.replace(/^\s+|\s+$/g, '').toLowerCase();
        const ok = { alt: "Alt", ctrl: "Ctrl", none: "None" }[normalized];
        if (ok) {
            modifierKey = ok;
            GM_setValue('or_modifierKey', ok);
            alert("Modifier set to: " + ok);
        } else {
            alert("Invalid. Must be Alt, Ctrl or None.");
        }
    });
    GM_registerMenuCommand("Toggle Navigation Panel", () => {
        panelEnabled = !panelEnabled;
        GM_setValue('or_panelEnabled', panelEnabled);
        if (panelEnabled) showPanel();
        else clearPanel();
        alert("Panel: " + (panelEnabled ? "Enabled" : "Disabled"));
    });

    /********************* SELECTORS/STRUCTURE ******************/
    function findScrollContainer() {
        return document.querySelector('main div.overflow-y-scroll') ||
               document.querySelector('main div[style*="overflow-y: auto;"]') ||
               document.querySelector('main div[style*="overflow-y: scroll;"]') ||
               document.querySelector('main');
    }

    function findMessageContainers() {
        if (!scrollContainer) return [];
        return Array.from(
            scrollContainer.querySelectorAll('div.duration-200.group.my-2.flex.flex-col.gap-2.md\\:my-0')
        ).filter(d => d.offsetParent !== null && d.querySelector('.flex.gap-2.items-center, .flex.gap-2.flex-row-reverse'));
    }

    function msgContentElem(msgDiv) {
        return msgDiv.querySelector('.overflow-auto') || msgDiv.querySelector('div.flex.max-w-full.flex-col.relative.overflow-auto');
    }

    function msgToggleExpandBtn(msgDiv) {
        return msgDiv.querySelector(
            'div.group.flex.flex-col.gap-2.items-start > div.flex.max-w-full.flex-col.relative.overflow-auto.gap-1.items-start.w-full > div > div > button, ' +
            'div.group.flex.flex-col.gap-2.items-end > div.flex.max-w-full.flex-col.relative.overflow-auto.gap-1.items-end.w-full > div > div > button'
        );
    }

    function msgHeader(msgDiv) {
        return msgDiv.querySelector('.group.flex.flex-col.gap-2.items-start > .flex.gap-2, .group.flex.flex-col.gap-2.items-end > .flex.gap-2') ||
               msgDiv.querySelector('.flex.gap-2.items-center, .flex.gap-2.flex-row-reverse');
    }

    function getSpeakerName(msgDiv) {
        const hdr = msgHeader(msgDiv);
        if (!hdr) return "";
        const a = hdr.querySelector('span a');
        if (a) return a.textContent.replace(/\|.*/,'').replace('(edited)','').trim();
        const span = hdr.querySelector('span');
        if (span) return span.textContent.replace(/\|.*/,'').replace('(edited)','').trim();
        return "";
    }

    function getSpeakerAvatar(msgDiv) {
        const hdr = msgHeader(msgDiv);
        if (!hdr) return "";
        const img = hdr.querySelector("picture img, img");
        if (img) return img.src;
        return "";
    }

    /*********** PANEL AND FLOATING UX COMPONENTS ************/
    let speakerElem = null, speakerImg = null, speakerName = null;
    function ensureSpeakerFloat() {
        if (document.querySelector("#openrouter-nav-speaker-float")) return;
        speakerElem = document.createElement("div");
        speakerElem.id = "openrouter-nav-speaker-float";
        speakerElem.innerHTML = `<img style="display:none"><span id="openrouter-speaker"></span>`;
        document.body.appendChild(speakerElem);
        speakerImg = speakerElem.querySelector('img');
        speakerName = speakerElem.querySelector('span');
        speakerElem.classList.remove('openrouter-nav-visible');
    }

    function showSpeaker(msgDiv) {
        if (!speakerElem) return;
        if (!msgDiv || !document.body.contains(msgDiv)) {
            speakerElem.classList.remove("openrouter-nav-visible");
            return;
        }
        const name = getSpeakerName(msgDiv);
        const imgSrc = getSpeakerAvatar(msgDiv);
        speakerName.textContent = name || "Speaker";
        if (imgSrc) { speakerImg.style.display = ""; speakerImg.src = imgSrc; }
        else { speakerImg.style.display = "none"; speakerImg.removeAttribute('src'); }
        speakerElem.classList.add("openrouter-nav-visible");
    }

    let panelElem = null;
    function showPanel() {
        clearPanel();
        if (!panelEnabled) return;
        panelElem = document.createElement("div");
        panelElem.id = "openrouter-nav-panel";
        panelElem.innerHTML = `
            <button class="openrouter-nav-btn" title="Previous Message (k)">
                <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
                    <polyline points="15 18 9 12 15 6"/>
                </svg>
            </button>
            <button class="openrouter-nav-btn" title="Next Message (j)">
                <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
                    <polyline points="9 18 15 12 9 6"/>
                </svg>
            </button>
            <span class="openrouter-nav-divider"></span>
            <button class="openrouter-nav-btn" title="Top (Home)">
                <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
                    <polyline points="18 15 12 9 6 15"/>
                </svg>
            </button>
            <button class="openrouter-nav-btn" title="Bottom (End)">
                <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
                    <polyline points="6 9 12 15 18 9"/>
                </svg>
            </button>
            <span class="openrouter-nav-divider"></span>
            <button class="openrouter-nav-btn" title="Expand/Collapse (l/h)">
                <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
                    <polyline points="7 13 12 18 17 13"/>
                    <polyline points="7 6 12 11 17 6"/>
                </svg>
            </button>
        `;
        document.body.appendChild(panelElem);
        const btns = panelElem.querySelectorAll('.openrouter-nav-btn');
        btns[0].onclick = () => navToMsg(-1);
        btns[1].onclick = () => navToMsg(1);
        btns[2].onclick = () => scrollMsgTop();
        btns[3].onclick = () => scrollMsgBottom();
        btns[4].onclick = () => toggleMsgExpand();
    }

    function clearPanel() {
        if (panelElem) { panelElem.remove(); panelElem = null; }
    }

    /******************** NAVIGATION LOGIC ********************/
    let scrollContainer = null;
    let allMessages = [], highlighted = null;
    let blockHighlightUntil = 0;
    let lastInteractedMsg = null;
    let latestInputEdit = 0, lastEditingMsg = null, editPasteProhibit = false;
    let collapseRestoreMsg = null;
    const ANTI_HYSTERESIS_MS = 50;
    const EDIT_LOCK_DURATION_MS = 3000;
    const COLLAPSE_SCROLL_LOCK_MS = 500;

    function updateMsgList() {
        let prevId = highlighted?.dataset?.ormsgid;
        allMessages = findMessageContainers();
        allMessages.forEach((m, i) => {
            if (!m.dataset.ormsgid) m.dataset.ormsgid = "msg-" + Date.now() + "-" + Math.random();
        });
        if (prevId) {
            highlighted = allMessages.find(m => m.dataset?.ormsgid === prevId);
        }
        if (!highlighted && allMessages.length > 0) {
            highlighted = allMessages[allMessages.length - 1];
        }
        allMessages.forEach(m => m.classList.toggle('openrouter-nav-highlight', m === highlighted));
        if (highlighted) {
            showSpeaker(highlighted);
        } else {
            showSpeaker(null);
        }
    }

    function highlightMsg(msgDiv, opts = {}) {
        if (!msgDiv || !document.body.contains(msgDiv)) return;
        if (editPasteProhibit && lastEditingMsg && lastEditingMsg !== msgDiv) return;
        if (Date.now() < blockHighlightUntil && !opts.force) return;
        if (highlighted) highlighted.classList.remove('openrouter-nav-highlight');
        highlighted = msgDiv;
        highlighted.classList.add('openrouter-nav-highlight');
        showSpeaker(highlighted);
        lastInteractedMsg = highlighted;
        if (opts.scrollIntoView) {
            highlighted.scrollIntoView({ behavior: "smooth", block: opts.block || "center" });
            if (opts.scrollTop) {
                let ct = msgContentElem(highlighted);
                if (ct) ct.scrollTop = 0;
            }
            if (opts.scrollBottom) {
                let ct = msgContentElem(highlighted);
                if (ct) ct.scrollTop = ct.scrollHeight;
            }
        }
    }

    function navToMsg(dir = 1) {
        if (!allMessages.length) return;
        let idx = highlighted ? allMessages.indexOf(highlighted) : -1;
        let nextIdx = idx + dir;
        if (nextIdx < 0) nextIdx = 0;
        if (nextIdx > allMessages.length - 1) nextIdx = allMessages.length - 1;
        blockHighlightUntil = Date.now() + 350;
        if (allMessages[nextIdx]) highlightMsg(allMessages[nextIdx], { scrollIntoView: true, force: true });
    }

    function scrollMsgTop() {
        if (!highlighted) return;
        let ct = msgContentElem(highlighted);
        if (ct) ct.scrollTop = 0;
        highlighted.scrollIntoView({ behavior: "smooth", block: "start" });
        blockHighlightUntil = Date.now() + 300;
    }

    function scrollMsgBottom() {
        if (!highlighted) return;
        let ct = msgContentElem(highlighted);
        if (ct) ct.scrollTop = ct.scrollHeight;
        highlighted.scrollIntoView({ behavior: "smooth", block: "end" });
        blockHighlightUntil = Date.now() + 300;
    }

    function toggleMsgExpand() {
        if (!highlighted) return;
        const btn = msgToggleExpandBtn(highlighted);
        if (!btn) return;
        handleToggleScroll(highlighted);
        btn.click();
    }

    function handleToggleScroll(msgDiv) {
        collapseRestoreMsg = msgDiv;
        const scrollContainer = findScrollContainer();
        const scrollTopBefore = scrollContainer.scrollTop;
        const msgTopBefore = msgDiv.offsetTop;
        const visualTop = msgTopBefore - scrollTopBefore;
        setTimeout(() => {
            let msg = allMessages.find(m => m.dataset.ormsgid === collapseRestoreMsg.dataset.ormsgid);
            if (msg) {
                const msgTopAfter = msg.offsetTop;
                scrollContainer.scrollTop = msgTopAfter - visualTop;
                highlightMsg(msg, { force: true });
                ensureScrollInBounds(msg);
            }
            collapseRestoreMsg = null;
            blockHighlightUntil = Date.now() + COLLAPSE_SCROLL_LOCK_MS;
        }, 210);
    }

    function refreshActiveMsg() {
        if (!highlighted) return;
        const refreshSvg = highlighted.querySelector('svg path[d="M16.023 9.348h4.992v-.001M2.985 19.644v-4.992m0 0h4.992m-4.993 0 3.181 3.183a8.25 8.25 0 0 0 13.803-3.7M4.031 9.865a8.25 8.25 0 0 1 13.803-3.7l3.181 3.182m0-4.991v4.99"]');
        if (refreshSvg) {
            refreshSvg.closest('button').click();
        }
    }

    function updateHighlightOnScroll() {
        if (Date.now() < blockHighlightUntil) return;
        if (editPasteProhibit && lastEditingMsg) {
            ensureScrollInBounds(lastEditingMsg);
            return;
        }
        let best = null, maxVH = 0;
        const containerRect = scrollContainer.getBoundingClientRect();
        allMessages.forEach(m => {
            const rect = m.getBoundingClientRect();
            let top = Math.max(rect.top, containerRect.top);
            let bot = Math.min(rect.bottom, containerRect.bottom);
            let visH = Math.max(0, bot - top);
            if (visH > maxVH && visH > 48) {
                maxVH = visH;
                best = m;
            }
        });
        if (best && best !== highlighted) {
            highlightMsg(best);
        }
    }

    /*************** EDITINGS/PASTE SCROLL BOUNDING ***********/
    let scrollLockTimeout = null;
    function enforceScrollBoundOnEdit() {
        const act = document.activeElement;
        if (act && act.closest('.duration-200.group.my-2.flex.flex-col.gap-2.md\\:my-0') && (act.matches('input:not([type="checkbox"]):not([type="radio"]), textarea, [contenteditable="true"]'))) {
            const activeMsg = act.closest('.duration-200.group.my-2.flex.flex-col.gap-2.md\\:my-0');
            if (activeMsg && document.body.contains(activeMsg)) {
                lastEditingMsg = activeMsg;
                latestInputEdit = Date.now();
                editPasteProhibit = true;
                highlightMsg(activeMsg, { force: true });
                ensureScrollInBounds(activeMsg);
                if (scrollLockTimeout) clearTimeout(scrollLockTimeout);
                scrollLockTimeout = setTimeout(() => {
                    if (Date.now() - latestInputEdit >= EDIT_LOCK_DURATION_MS) {
                        editPasteProhibit = false;
                        lastEditingMsg = null;
                        scrollLockTimeout = null;
                    }
                }, EDIT_LOCK_DURATION_MS);
            }
        }
    }

    function ensureScrollInBounds(msgDiv) {
        if (!msgDiv || !scrollContainer) return;
        const msgRect = msgDiv.getBoundingClientRect();
        const scRect = scrollContainer.getBoundingClientRect();
        if (msgRect.top < scRect.top || msgRect.bottom > scRect.bottom) {
            msgDiv.scrollIntoView({ behavior: "auto", block: "center" });
        }
    }

    function disableContainerScroll() {
        if (scrollContainer) scrollContainer.style.overflowY = 'hidden';
    }

    function enableContainerScroll() {
        if (scrollContainer) scrollContainer.style.overflowY = 'auto';
    }

    /*************** LISTENERS: PANEL+CLICK **************/
    function panelAndPageListeners() {
        scrollContainer.addEventListener('click', e => {
            const msg = e.target.closest('.duration-200.group.my-2.flex.flex-col.gap-2.md\\:my-0');
            if (msg && allMessages.includes(msg)) highlightMsg(msg, { force: true });
        });
        scrollContainer.addEventListener('focusin', e => {
            const msg = e.target.closest('.duration-200.group.my-2.flex.flex-col.gap-2.md\\:my-0, [data-ormsgid]');
            if (msg && allMessages.includes(msg)) highlightMsg(msg, { force: true });
        });
        scrollContainer.addEventListener('mousedown', e => {
            const msg = e.target.closest('.duration-200.group.my-2.flex.flex-col.gap-2.md\\:my-0, [data-ormsgid]');
            if (msg && allMessages.includes(msg)) highlightMsg(msg, { force: true });
        });
        const observer = new MutationObserver(() => {
            requestAnimationFrame(() => {
                requestAnimationFrame(updateMsgList);
            });
        });
        observer.observe(scrollContainer, { childList: true, subtree: true });

        const expandCollapseSelector = 'div.group.flex.flex-col.gap-2.items-start > div.flex.max-w-full.flex-col.relative.overflow-auto.gap-1.items-start.w-full > div > div > button, ' +
                                       'div.group.flex.flex-col.gap-2.items-end > div.flex.max-w-full.flex-col.relative.overflow-auto.gap-1.items-end.w-full > div > div > button';
        scrollContainer.addEventListener('mousedown', e => {
            const btn = e.target.closest(expandCollapseSelector);
            if (btn) {
                const msgDiv = btn.closest('.duration-200.group.my-2.flex.flex-col.gap-2.md\\:my-0');
                if (msgDiv && allMessages.includes(msgDiv)) {
                    handleToggleScroll(msgDiv);
                }
            }
        });
    }

    /***************** HOTKEYS (NAV/CTRL) ********************/
    function isModifier(event) {
        if (modifierKey === "None") return !event.ctrlKey && !event.altKey;
        if (modifierKey === "Ctrl") return event.ctrlKey && !event.altKey;
        if (modifierKey === "Alt") return event.altKey && !event.ctrlKey;
        return false;
    }

    function setupHotkeys() {
        document.addEventListener('keydown', function (e) {
            if (
                e.target.matches('input, textarea, [contenteditable]') &&
                !["Home", "End", "PageUp", "PageDown"].includes(e.key)
            ) return;
            if (!isModifier(e)) return;
            let handled = false;
            switch (e.key) {
                case 'j':
                    navToMsg(1);
                    handled = true;
                    break;
                case 'k':
                    navToMsg(-1);
                    handled = true;
                    break;
                case 'l':
                case 'h':
                    toggleMsgExpand();
                    handled = true;
                    break;
                case 'Home':
                    if (!e.target.matches('[contenteditable]')) {
                        scrollMsgTop();
                        handled = true;
                    }
                    break;
                case 'End':
                    if (!e.target.matches('[contenteditable]')) {
                        scrollMsgBottom();
                        handled = true;
                    }
                    break;
                case 'r':
                    refreshActiveMsg();
                    handled = true;
                    break;
            }
            if (handled) e.preventDefault();
        });
    }

    /******************** INIT ENTRYPOINT ********************/
    async function initPowerNav() {
        const waitFor = (f) => new Promise(resolve => {
            function step() {
                const x = f();
                if (x) resolve(x);
                else setTimeout(step, 220);
            }
            step();
        });
        scrollContainer = await waitFor(findScrollContainer);

        ensureSpeakerFloat();
        if (panelEnabled) showPanel();

        updateMsgList();
        panelAndPageListeners();
        setupHotkeys();

        let lastScrollUpd = 0;
        scrollContainer.addEventListener('scroll', () => {
            if (Date.now() - lastScrollUpd > ANTI_HYSTERESIS_MS) {
                updateHighlightOnScroll();
                lastScrollUpd = Date.now();
            }
            if (editPasteProhibit && lastEditingMsg) {
                ensureScrollInBounds(lastEditingMsg);
            }
            const active = document.activeElement;
            if (active && (active.matches('input, textarea, [contenteditable]'))) {
                active.blur();
            }
        }, { passive: true });

        document.addEventListener('input', enforceScrollBoundOnEdit, true);
        document.addEventListener('paste', (e) => {
            enforceScrollBoundOnEdit();
            disableContainerScroll();
            setTimeout(enableContainerScroll, 100);
        }, true);
        document.addEventListener('cut', enforceScrollBoundOnEdit, true);

        document.addEventListener('focusout', () => {
            if (editPasteProhibit && Date.now() - latestInputEdit > EDIT_LOCK_DURATION_MS / 2) {
                editPasteProhibit = false;
                lastEditingMsg = null;
            }
        }, true);

        setInterval(updateMsgList, 880);

        document.addEventListener('visibilitychange', () => {
            if (document.visibilityState === "visible") setTimeout(updateMsgList, 500);
        });

        window.addEventListener('resize', () => { setTimeout(updateHighlightOnScroll, 80); });
    }

    initPowerNav();
})();