import { app } from "/scripts/app.js";
import { api } from "/scripts/api.js";
import { ComfyWidgets } from "/scripts/widgets.js";

const CHAT_CLASSES = ["ChatNode", "Chat"];
const VIEWER_CLASS = "ChatHistoryViewer";
const ROUTE_CHAT = "/chat_optimize/chat";
let historyModal = null;

function ensureHistoryModal() {
    if (historyModal) return historyModal;

    const overlay = document.createElement("div");
    overlay.id = "chat-optimize-history-modal";
    overlay.style.position = "fixed";
    overlay.style.top = "0";
    overlay.style.left = "0";
    overlay.style.width = "100%";
    overlay.style.height = "100%";
    overlay.style.display = "none";
    overlay.style.alignItems = "center";
    overlay.style.justifyContent = "center";
    overlay.style.background = "rgba(0, 0, 0, 0.6)";
    overlay.style.zIndex = "9999";

    const panel = document.createElement("div");
    panel.style.width = "min(900px, 92%)";
    panel.style.height = "80%";
    panel.style.background = "#111";
    panel.style.color = "#ddd";
    panel.style.border = "1px solid #333";
    panel.style.borderRadius = "8px";
    panel.style.display = "flex";
    panel.style.flexDirection = "column";
    panel.style.boxShadow = "0 10px 30px rgba(0, 0, 0, 0.5)";

    const header = document.createElement("div");
    header.style.display = "flex";
    header.style.alignItems = "center";
    header.style.justifyContent = "space-between";
    header.style.padding = "10px 12px";
    header.style.borderBottom = "1px solid #333";
    header.style.background = "#151515";

    const title = document.createElement("div");
    title.textContent = "Chat History";
    title.style.fontSize = "14px";
    title.style.fontWeight = "600";

    const closeBtn = document.createElement("button");
    closeBtn.textContent = "Close";
    closeBtn.style.background = "#222";
    closeBtn.style.color = "#ddd";
    closeBtn.style.border = "1px solid #444";
    closeBtn.style.borderRadius = "6px";
    closeBtn.style.padding = "6px 10px";
    closeBtn.style.cursor = "pointer";
    closeBtn.addEventListener("click", () => closeHistoryModal());

    header.appendChild(title);
    header.appendChild(closeBtn);

    const content = document.createElement("pre");
    content.style.flex = "1";
    content.style.margin = "0";
    content.style.padding = "12px";
    content.style.overflow = "auto";
    content.style.whiteSpace = "pre-wrap";
    content.style.fontFamily = "monospace";
    content.style.fontSize = "13px";

    panel.appendChild(header);
    panel.appendChild(content);
    overlay.appendChild(panel);
    overlay.addEventListener("click", (e) => {
        if (e.target === overlay) closeHistoryModal();
    });

    document.body.appendChild(overlay);
    const onKey = (e) => {
        if (e.key === "Escape") closeHistoryModal();
    };
    document.addEventListener("keydown", onKey);

    historyModal = { overlay, content, onKey };
    return historyModal;
}

function openHistoryModal(text) {
    const modal = ensureHistoryModal();
    modal.content.textContent = text || "";
    modal.overlay.style.display = "flex";
    setTimeout(() => {
        modal.content.scrollTop = modal.content.scrollHeight;
    }, 0);
}

function closeHistoryModal() {
    if (!historyModal) return;
    historyModal.overlay.style.display = "none";
}

function isHistoryModalOpen() {
    return historyModal && historyModal.overlay.style.display !== "none";
}

function setWidgetValue(node, name, value) {
    if (!node.widgets) return;
    const w = node.widgets.find((w) => w.name === name);
    if (w) {
        w.value = value ?? "";
        if (w.onChange) w.onChange(value);
        node.setDirtyCanvas(true, true);
    }
}

function addOrGetWidget(node, type, name, opts = {}) {
    if (!node.widgets) node.widgets = [];
    let w = node.widgets.find((w) => w.name === name);
    if (w) return w;
    w = node.addWidget(type, name, opts.default ?? "", () => { }, opts);
    return w;
}

async function postJSON(path, body) {
    const res = await api.fetchApi(path, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(body),
    });
    if (!res.ok) throw new Error(await res.text());
    return await res.json();
}

function propagateOutput(node, slotIndex, value) {
    if (!node.outputs || !node.outputs[slotIndex] || !node.outputs[slotIndex].links) return;

    node.outputs[slotIndex].links.forEach((linkId) => {
        const link = app.graph.links[linkId];
        if (!link) return;
        const targetNode = app.graph.getNodeById(link.target_id);
        if (!targetNode) return;

        if (targetNode.inputs && targetNode.inputs[link.target_slot] && targetNode.inputs[link.target_slot].widget) {
            const widgetName = targetNode.inputs[link.target_slot].widget.name;
            const w = targetNode.widgets?.find(w => w.name === widgetName);
            if (w) {
                w.value = value;
                if (w.onChange) w.onChange(value);
                targetNode.setDirtyCanvas(true, true);
                return;
            }
        }
    });
}

function getLLMConfig(node) {
    const inputIndex = node.inputs?.findIndex(i => i.name === "llm_config");
    if (inputIndex === undefined || inputIndex === -1) return null;

    const input = node.inputs[inputIndex];
    if (!input.link) return null;

    const linkId = input.link;
    const link = app.graph.links[linkId];
    if (!link) return null;

    const sourceNode = app.graph.getNodeById(link.origin_id);
    if (!sourceNode) return null;

    const config = {};
    const providerW = sourceNode.widgets?.find(w => w.name === "provider");
    if (providerW) config.provider = providerW.value;

    const baseUrlW = sourceNode.widgets?.find(w => w.name === "base_url");
    if (baseUrlW) config.base_url = baseUrlW.value;

    const modelW = sourceNode.widgets?.find(w => w.name === "model_name");
    if (modelW) config.model_name = modelW.value;

    const tempW = sourceNode.widgets?.find(w => w.name === "temperature");
    if (tempW) config.temperature = tempW.value;

    const topPW = sourceNode.widgets?.find(w => w.name === "top_p");
    if (topPW) config.top_p = topPW.value;

    const maxTokensW = sourceNode.widgets?.find(w => w.name === "max_new_tokens");
    if (maxTokensW) config.max_new_tokens = maxTokensW.value;

    const tokenW = sourceNode.widgets?.find(w => w.name === "hf_token");
    if (tokenW) config.hf_token = tokenW.value;

    const hfApiUrlW = sourceNode.widgets?.find(w => w.name === "hf_api_url");
    if (hfApiUrlW) config.hf_api_url = hfApiUrlW.value;

    const apiKeyW = sourceNode.widgets?.find(w => w.name === "api_key");
    if (apiKeyW) config.api_key = apiKeyW.value;

    return config;
}

function updateWidgetLock(node) {
    const input = node.inputs?.find(i => i.name === "llm_config");
    const isConnected = input && input.link !== null;

    ["base_url", "model_name"].forEach(name => {
        const w = node.widgets?.find(w => w.name === name);
        if (w) {
            w.disabled = isConnected;
        }
    });
}

function setupChatNode(node) {
    console.log("[ChatNode] Setting up");
    
    let statusWidget = addOrGetWidget(node, "text", "status", { default: "idle" });
    if (statusWidget.inputEl) {
        statusWidget.inputEl.readOnly = true;
        statusWidget.inputEl.style.fontSize = "11px";
        statusWidget.inputEl.style.color = "#888";
        statusWidget.inputEl.style.backgroundColor = "#1a1a1a";
        statusWidget.inputEl.style.border = "1px solid #333";
    }

    let historyStore = addOrGetWidget(node, "hidden", "history_preview", { default: "" });
    let responseStore = addOrGetWidget(node, "hidden", "assistant_response_preview", { default: "" });

    const sendAction = async () => {
        const actionWidget = node.widgets?.find((w) => w.name === "action");
        const action = actionWidget ? actionWidget.value : "send";
        
        try {
            statusWidget.value = `working: ${action}...`;
            node.setDirtyCanvas(true, true);
            
            const payload = {};
            const excludeWidgets = ["status", "history_preview", "assistant_response_preview", "auto_clear_input"];
            
            (node.widgets || []).forEach((w) => {
                if (excludeWidgets.includes(w.name)) return;
                payload[w.name] = w.value;
            });
            payload["action"] = action;

            const llmConfig = getLLMConfig(node);
            if (llmConfig) payload["llm_config"] = llmConfig;

            const data = await postJSON(ROUTE_CHAT, payload);
            
            responseStore.value = data.assistant_response || "";
            historyStore.value = data.readable_history || "";

            propagateOutput(node, 0, data.assistant_response || "");
            propagateOutput(node, 1, data.readable_history || "");

            const autoClear = node.widgets?.find((w) => w.name === "auto_clear_input")?.value;
            if (action === "send" && autoClear) {
                setWidgetValue(node, "user_message", "");
            }

            statusWidget.value = "done";
            node.setDirtyCanvas(true, true);
            
            console.log("[ChatNode] Response received");
        } catch (err) {
            console.error("[ChatNode] Error:", err);
            statusWidget.value = `error: ${err.message}`;
            node.setDirtyCanvas(true, true);
        }
    };

    node.onConnectionsChange = () => setTimeout(() => updateWidgetLock(node), 100);
    setTimeout(() => updateWidgetLock(node), 100);

    const existingButton = node.widgets?.find(w => w.name === "Execute");
    if (!existingButton) {
        node.addWidget("button", "Execute", null, () => sendAction());
    }
}

function setupHistoryViewer(node) {
    console.log("[ChatHistoryViewer] Setting up stable scroll version");
    
    // 1. Find or create the history widget
    let historyWidget = node.widgets?.find(w => w.name === "history");
    if (!historyWidget) {
        historyWidget = node.addWidget("text", "history", "", () => {}, { multiline: true });
    }
    const getHistoryText = () => historyWidget.value ?? historyWidget.inputEl?.value ?? "";

    const existingButton = node.widgets?.find(w => w.name === "Open History");
    if (!existingButton) {
        node.addWidget("button", "Open History", null, () => openHistoryModal(getHistoryText()));
    }

    // 2. Clear default drawing of text on canvas
    node.onDrawForeground = function(ctx) {
        // Stop default text preview if any
    };

    // 3. Configure the textarea
    if (historyWidget.inputEl) {
        const el = historyWidget.inputEl;
        el.readOnly = true;
        el.style.display = "block"; // Force visible
        el.style.visibility = "visible";
        el.style.pointerEvents = "auto"; // Ensure scrollable
        el.style.height = "100%";
        el.style.width = "100%";
        el.style.backgroundColor = "#121212";
        el.style.color = "#ddd";
        el.style.fontFamily = "monospace";
        el.style.fontSize = "13px";
        el.style.padding = "10px";
        el.style.border = "none";
        el.style.boxSizing = "border-box";
        el.style.resize = "none";
        el.style.overflowY = "scroll";
        el.style.whiteSpace = "pre-wrap";
        
        // Custom scrollbar
        el.parentElement.style.overflow = "hidden"; // Container shouldn't scroll
    }

    // 4. Override widget visibility when connected
    // ComfyUI usually hides widgets linked to inputs. We want this one visible.
    const originalUpdate = node.onNodeCreated;
    historyWidget.type = "text"; // Keep it as text so it renders the textarea

    // Prevent ComfyUI from hiding it
    const originalOnConnectionsChange = node.onConnectionsChange;
    node.onConnectionsChange = function() {
        if (originalOnConnectionsChange) originalOnConnectionsChange.apply(this, arguments);
        if (historyWidget.inputEl) {
           historyWidget.inputEl.style.display = "block";
           historyWidget.inputEl.style.visibility = "visible";
        }
    };

    // 5. High-frequency sync and auto-scroll
    const originalOnChange = historyWidget.onChange;
    historyWidget.onChange = function(v) {
        if (originalOnChange) originalOnChange.call(this, v);
        // Sync DOM explicitly if needed
        if (historyWidget.inputEl && historyWidget.inputEl.value !== v) {
            historyWidget.inputEl.value = v || "";
        }
        if (isHistoryModalOpen()) {
            openHistoryModal(v || "");
        }
        // Auto scroll to bottom
        setTimeout(() => {
            if (historyWidget.inputEl) {
                historyWidget.inputEl.scrollTop = historyWidget.inputEl.scrollHeight;
            }
        }, 10);
    };

    // 6. Set initial size
    node.setSize([600, 500]);
}

app.registerExtension({
    name: "chat_optimize_ui",
    priority: 10,
    async nodeCreated(node) {
        if (CHAT_CLASSES.includes(node.comfyClass)) {
            setupChatNode(node);
        } else if (node.comfyClass === VIEWER_CLASS) {
            setupHistoryViewer(node);
        }
    },
});