// ==UserScript==
// @name         Grok Imagine Pro UI v2.5
// @namespace    http://tampermonkey.net/
// @version      2.5
// @description  Advanced UI for x.ai with Prompt History and Gallery Management
// @author       Gemini-3-flash
// @match        https://console.x.ai/team/*/imagine
// @grant        none
// ==/UserScript==

(function() {
    'use strict';

    let historyImages = []; // Сгенерированные картинки
    let selectedImages = []; // Хранилище Base64 выбранных файлов
    let promptHistory = JSON.parse(localStorage.getItem('grok_prompt_history') || '[]');
    let currentIndex = -1;

    const style = document.createElement('style');
    style.innerHTML = `
        #custom-gen-ui {
            position: fixed; top: 0; left: 0; width: 100%; height: 100%;
            background: #0a0a0a; color: #fff; z-index: 9998;
            display: flex; flex-direction: row; font-family: sans-serif;
        }

        /* SIDEBAR */
        #sidebar {
            width: 300px; background: #111; border-right: 1px solid #333;
            display: flex; flex-direction: column; transition: 0.3s; overflow: hidden;
        }
        #sidebar.collapsed { width: 0; border: none; }
        .sidebar-header { padding: 15px; font-weight: bold; border-bottom: 1px solid #222; display: flex; justify-content: space-between; align-items: center; white-space: nowrap; }
        #prompt-list { flex: 1; overflow-y: auto; padding: 10px; }
        .prompt-item {
            background: #1a1a1a; padding: 10px; border-radius: 6px; margin-bottom: 8px;
            font-size: 13px; cursor: pointer; position: relative; border: 1px solid transparent; transition: 0.2s;
        }
        .prompt-item:hover { border-color: #555; background: #222; }
        .prompt-item .delete-prompt { position: absolute; right: 5px; top: 5px; color: #666; padding: 2px 5px; }
        .prompt-item .delete-prompt:hover { color: #ff4444; }

        /* MAIN AREA */
        #main-content { flex: 1; display: flex; flex-direction: column; padding: 20px; overflow-y: auto; position: relative; }
        .controls { width: 100%; max-width: 1000px; margin: 0 auto; display: flex; flex-direction: column; gap: 12px; background: #161616; padding: 20px; border-radius: 12px; border: 1px solid #333; }
        textarea { width: 100%; height: 280px; background: #000; color: #eee; border: 1px solid #444; padding: 12px; border-radius: 8px; font-size: 15px; }
        .row { display: flex; gap: 12px; align-items: center; flex-wrap: wrap; }

        select, input[type="number"] { background: #000; color: #fff; border: 1px solid #444; padding: 8px; border-radius: 6px; }
        button { border: none; padding: 10px 18px; font-weight: bold; cursor: pointer; border-radius: 6px; transition: 0.2s; }
        #gen-btn { background: #fff; color: #000; }
        #gen-btn:hover { background: #ccc; }
        #gen-btn:disabled { background: #444; color: #888; }
        .btn-secondary { background: #333; color: #fff; }
        .btn-secondary:hover { background: #444; }
        .btn-danger { background: #422; color: #f66; border: 1px solid #633; }
        .btn-danger:hover { background: #622; }

        #results { margin-top: 25px; display: grid; grid-template-columns: repeat(auto-fill, minmax(220px, 1fr)); gap: 15px; width: 100%; }
        .img-card { background: #161616; border: 1px solid #333; padding: 5px; border-radius: 8px; cursor: pointer; position: relative; }
        .img-card img { width: 100%; height: 200px; object-fit: cover; border-radius: 4px; }

        /* MODAL */
        #img-modal { position: fixed; top: 0; left: 0; width: 100%; height: 100%; background: rgba(0,0,0,0.95); z-index: 10000; display: none; align-items: center; justify-content: center; }
        #modal-img { max-width: 90%; max-height: 90vh; border-radius: 4px; }
        .modal-btn { position: absolute; background: rgba(255,255,255,0.1); color: #fff; border: none; width: 50px; height: 50px; border-radius: 50%; font-size: 24px; cursor: pointer; display: flex; align-items: center; justify-content: center; }
        #modal-close { top: 20px; right: 20px; }
        #modal-prev { left: 20px; }
        #modal-next { right: 20px; }
        #modal-download { position:absolute;top:40px; left:auto;width:auto; padding:0 25px; background:var(--color-green-500); color:#000; text-decoration:none; line-height:44px; font-weight:bold; border-radius:6px; box-shadow:black 0px 0px 10px 2px }
        .status { color: #888; font-size: 13px; font-family: monospace; }

        /* Превью референсов */
        #reference-previews {
            display: flex; gap: 10px; margin-top: 10px; flex-wrap: wrap;
            min-height: 20px; padding: 10px; border: 2px dashed transparent;
            border-radius: 8px; transition: 0.2s;
        }
        #reference-previews.drag-over { border-color: #0f0; background: rgba(0,255,0,0.05); }

        .ref-item {
            width: 80px; height: 80px; position: relative; cursor: grab;
            border: 2px solid #333; border-radius: 6px; overflow: hidden;
            background: #000; transition: transform 0.2s;
        }
        .ref-item:active { cursor: grabbing; }
        .ref-item img { width: 100%; height: 100%; object-fit: cover; pointer-events: none; }
        .ref-item .remove-ref {
            position: absolute; top: 2px; right: 2px; background: rgba(0,0,0,0.7);
            color: #ff4444; border: none; border-radius: 4px; padding: 0 4px;
            font-size: 12px; cursor: pointer; z-index: 2;
        }
        .ref-item .ref-index {
            position: absolute; bottom: 2px; left: 2px; background: rgba(0,0,0,0.7);
            color: #fff; font-size: 10px; padding: 0 4px; border-radius: 4px;
        }
        .ref-item.dragging { opacity: 0.5; transform: scale(0.9); }
    `;
    document.head.appendChild(style);

    const sleep = (ms) => new Promise(res => setTimeout(res, ms));

    function createUI() {
        document.querySelectorAll('aside, main').forEach(el => el.remove());
        const ui = document.createElement('div');
        ui.id = 'custom-gen-ui';
        ui.innerHTML = `
            <div id="sidebar">
                <div class="sidebar-header">
                    <span>PROMPT HISTORY</span>
                    <button class="btn-secondary" id="close-sidebar" style="padding: 2px 8px;">✕</button>
                </div>
                <div id="prompt-list"></div>
            </div>
            <div id="main-content">
                <div class="controls">
                    <textarea id="prompt-input" placeholder="What should I imagine?"></textarea>
                    <div class="row">
                        <button class="btn-secondary" id="toggle-sidebar">📜 History</button>
                        <input type="file" id="file-input" multiple accept="image/*" style="display:none">
                        <button class="btn-secondary" id="upload-btn">📎 Edit (0/10)</button>
                        <button class="btn-danger" id="clear-files" style="display:none; padding: 2px 8px;">✕</button>
                        <select id="ratio-select">
                            <option value="auto">Auto Ratio</option>
                            <option value="1:1">1:1</option>
                            <option value="16:9">16:9</option>
                            <option value="9:16">9:16</option>
                            <option value="19.5:9">19.5:9</option>
                            <option value="9:19.5">9:19.5</option>
                            <option value="20:9">20:9</option>
                            <option value="9:20">9:20</option>
                            <option value="4:3">4:3</option>
                            <option value="3:2">3:2</option>
                            <option value="2:3">2:3</option>
                            <option value="2:1">2:1</option>
                            <option value="1:2">1:2</option>
                        </select>
                        <input type="number" id="count-input" value="1" min="1" max="20" style="width: 45px;">
                        <select id="mode-select">
                            <option value="parallel">Parallel</option>
                            <option value="native">Native (n)</option>
                            <option value="stepwise">Stepwise</option>
                        </select>
                        <button id="gen-btn">GENERATE</button>
                        <button class="btn-danger" id="clear-gallery">Clear Gallery</button>
                        <div id="reference-previews"></div>
                        <span id="global-status" class="status"></span>
                    </div>
                </div>
                <div id="results"></div>
            </div>

            <div id="img-modal">
                <button id="modal-close" class="modal-btn">✕</button>
                <button id="modal-prev" class="modal-btn">←</button>
                <img id="modal-img" src="">
                <button id="modal-next" class="modal-btn">→</button>
                <a id="modal-download" download="grok.png">Download</a>
            </div>
        `;
        document.body.appendChild(ui);

        // Bindings
        document.getElementById('gen-btn').onclick = startGeneration;
        document.getElementById('clear-gallery').onclick = clearGallery;
        document.getElementById('toggle-sidebar').onclick = toggleSidebar;
        document.getElementById('close-sidebar').onclick = toggleSidebar;

        document.getElementById('modal-close').onclick = closeModal;
        document.getElementById('modal-prev').onclick = () => navigate(-1);
        document.getElementById('modal-next').onclick = () => navigate(1);
        document.getElementById('img-modal').onclick = (e) => { if(e.target.id === 'img-modal') closeModal(); };

        document.addEventListener('keydown', (e) => {
            if (document.getElementById('img-modal').style.display === 'flex') {
                if (e.key === 'Escape') closeModal();
                if (e.key === 'ArrowLeft') navigate(-1);
                if (e.key === 'ArrowRight') navigate(1);
            }
        });

        const refContainer = document.getElementById('reference-previews');
        const fileInput = document.getElementById('file-input');
        const uploadBtn = document.getElementById('upload-btn');
        const ratioSelect = document.getElementById('ratio-select');
        const genBtn = document.getElementById('gen-btn');

        // Функция обновления UI референсов
        function renderRefPreviews() {
            refContainer.innerHTML = '';
            selectedImages.forEach((imgObj, index) => {
                const item = document.createElement('div');
                item.className = 'ref-item';
                item.draggable = true;
                item.dataset.index = index;
                item.innerHTML = `
                <img src="${imgObj.url}">
                <button class="remove-ref">✕</button>
                <span class="ref-index">${index + 1}</span>
            `;

                // Удаление одного референса
                item.querySelector('.remove-ref').onclick = (e) => {
                    e.stopPropagation();
                    selectedImages.splice(index, 1);
                    updateEditState();
                };

                // Логика сортировки (Drag & Drop элементов)
                item.ondragstart = (e) => {
                    e.dataTransfer.setData('text/plain', index);
                    item.classList.add('dragging');
                };
                item.ondragend = () => item.classList.remove('dragging');

                item.ondragover = (e) => e.preventDefault();
                item.ondrop = (e) => {
                    e.preventDefault();
                    const fromIdx = parseInt(e.dataTransfer.getData('text/plain'));
                    const toIdx = index;
                    if (fromIdx === toIdx) return;

                    const movedItem = selectedImages.splice(fromIdx, 1)[0];
                    selectedImages.splice(toIdx, 0, movedItem);
                    updateEditState();
                };

                refContainer.appendChild(item);
            });
        }

        function updateEditState() {
            const count = selectedImages.length;
            uploadBtn.innerText = `📎 Edit (${count}/10)`;
            uploadBtn.style.color = count > 0 ? "#0f0" : "#fff";
            ratioSelect.disabled = count > 0;
            if (count > 0) ratioSelect.value = "auto";
            genBtn.innerText = count > 0 ? "EDIT IMAGES" : "GENERATE";
            renderRefPreviews();
        }

        // Обработка добавления файлов
        async function handleFiles(files) {
            const remainingSlots = 10 - selectedImages.length;
            const filesToAdd = Array.from(files).slice(0, remainingSlots);

            for (const file of filesToAdd) {
                const base64 = await new Promise(resolve => {
                    const reader = new FileReader();
                    reader.onload = () => resolve(reader.result);
                    reader.readAsDataURL(file);
                });
                selectedImages.push({ url: base64 });
            }
            updateEditState();
        }

        // События для кнопки и инпута
        uploadBtn.onclick = () => fileInput.click();
        fileInput.onchange = (e) => handleFiles(e.target.files);

        // Drag-and-drop файлов на контейнер или кнопку
        [uploadBtn, refContainer].forEach(el => {
            el.ondragover = (e) => {
                e.preventDefault();
                refContainer.classList.add('drag-over');
            };
            el.ondragleave = () => refContainer.classList.remove('drag-over');
            el.ondrop = (e) => {
                e.preventDefault();
                refContainer.classList.remove('drag-over');
                if (e.dataTransfer.files.length > 0) {
                    handleFiles(e.dataTransfer.files);
                }
            };
        });

        renderHistory();
    }

    // --- HISTORY LOGIC ---
    function savePrompt(text) {
        if (!text || promptHistory[0] === text) return;
        promptHistory = [text, ...promptHistory.filter(t => t !== text)].slice(0, 50);
        localStorage.setItem('grok_prompt_history', JSON.stringify(promptHistory));
        renderHistory();
    }

    function renderHistory() {
        const list = document.getElementById('prompt-list');
        list.innerHTML = '';
        promptHistory.forEach((text, i) => {
            const item = document.createElement('div');
            item.className = 'prompt-item';
            item.title = text;
            item.innerHTML = `
                <div class="text-truncate">${text.substring(0, 80)}${text.length > 80 ? '...' : ''}</div>
                <span class="delete-prompt" data-idx="${i}">✕</span>
            `;
            item.onclick = (e) => {
                if (e.target.className !== 'delete-prompt') {
                    document.getElementById('prompt-input').value = text;
                }
            };
            item.querySelector('.delete-prompt').onclick = (e) => {
                e.stopPropagation();
                promptHistory.splice(i, 1);
                localStorage.setItem('grok_prompt_history', JSON.stringify(promptHistory));
                renderHistory();
            };
            list.appendChild(item);
        });
    }

    function toggleSidebar() {
        document.getElementById('sidebar').classList.toggle('collapsed');
    }

    function clearGallery() {
        if (confirm("Clear all images from current gallery?")) {
            historyImages = [];
            document.getElementById('results').innerHTML = '';
            currentIndex = -1;
        }
    }

    // --- REQUEST LOGIC ---
    async function sendRequest(prompt, ratio, n = 1) {
        const isEdit = selectedImages.length > 0;
        const url = isEdit ? 'https://console.x.ai/v1/images/edits' : 'https://console.x.ai/v1/images/generations';

        const body = {
            model: 'grok-imagine-image',
            prompt, n,
            aspect_ratio: isEdit ? 'auto' : ratio,
            resolution: '2k',
            response_format: 'b64_json'
        };

        if (isEdit) {
            body.images = selectedImages; // Добавляем массив картинок
        }

        const response = await fetch(url, {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify(body)
        });

        if (!response.ok) {
            const err = await response.json();
            throw new Error(err.error || `HTTP ${response.status}`);
        }
        return await response.json();
    }

    function addImageToGallery(b64Data) {
        const dataUrl = `data:image/png;base64,${b64Data}`;
        const idx = historyImages.length;
        historyImages.push({ url: dataUrl, name: `grok_${Date.now()}.png` });

        const card = document.createElement('div');
        card.className = 'img-card';
        card.innerHTML = `<img src="${dataUrl}">`;
        card.onclick = () => openModal(idx);
        document.getElementById('results').prepend(card);
    }

    async function startGeneration() {
        const prompt = document.getElementById('prompt-input').value;
        const ratio = document.getElementById('ratio-select').value;
        const count = parseInt(document.getElementById('count-input').value);
        const mode = document.getElementById('mode-select').value;
        const status = document.getElementById('global-status');
        const btn = document.getElementById('gen-btn');

        if (!prompt) return;
        btn.disabled = true;
        savePrompt(prompt);

        try {
            if (mode === 'native') {
                status.innerText = "Native batching...";
                const res = await sendRequest(prompt, ratio, count);
                res.data.forEach(item => addImageToGallery(item.b64_json));
            } else if (mode === 'parallel') {
                status.innerText = `Parallel: 0/${count}`;
                let done = 0;
                const tasks = Array.from({ length: count }).map(() =>
                    sendRequest(prompt, ratio, 1).then(r => {
                        addImageToGallery(r.data[0].b64_json);
                        done++;
                        status.innerText = `Parallel: ${done}/${count}`;
                    })
                );
                await Promise.all(tasks);
            } else if (mode === 'stepwise') {
                for (let i = 0; i < count; i++) {
                    status.innerText = `Stepwise: ${i+1}/${count}`;
                    const res = await sendRequest(prompt, ratio, 1);
                    addImageToGallery(res.data[0].b64_json);
                    if (i < count - 1) {
                        const wait = Math.floor(Math.random() * 8000) + 5000;
                        for(let s = wait/1000; s > 0; s--) {
                            status.innerText = `Wait ${Math.round(s)}s (${i+1}/${count})`;
                            await sleep(1000);
                        }
                    }
                }
            }
            status.innerText = "Done.";
        } catch (e) {
            status.innerText = `Error: ${e.message}`;
        } finally {
            btn.disabled = false;
        }
    }

    // --- MODAL & NAV ---
    function openModal(i) {
        currentIndex = i;
        document.getElementById('modal-img').src = historyImages[i].url;
        document.getElementById('modal-download').href = historyImages[i].url;
        document.getElementById('modal-download').download = historyImages[i].name;
        document.getElementById('img-modal').style.display = 'flex';
        document.body.style.overflow = 'hidden';
    }
    function closeModal() {
        document.getElementById('img-modal').style.display = 'none';
        document.body.style.overflow = 'auto';
    }
    function navigate(dir) {
        currentIndex = (currentIndex + dir + historyImages.length) % historyImages.length;
        openModal(currentIndex);
    }

    const checkExist = setInterval(() => {
        if (document.querySelector('textarea, [role="textbox"]')) {
            clearInterval(checkExist);
            createUI();
        }
    }, 1000);
})();