let sessions = {};
let activeSessionId = null;
let activeAbortController = null;

// Initialize on DOM load
document.addEventListener('DOMContentLoaded', () => {
    loadSessions();

    const textarea = document.getElementById('userInput');
    if (textarea) {
        textarea.addEventListener('input', function() {
            this.style.height = 'auto';
            this.style.height = (this.scrollHeight) + 'px';
            if (this.scrollHeight > 200) {
                this.style.overflowY = 'auto';
            } else {
                this.style.overflowY = 'hidden';
            }
        });
        textarea.addEventListener('keydown', function(e) {
            if (e.key === 'Enter' && !e.shiftKey) {
                e.preventDefault();
                sendMessage();
            }
        });
    }
});

function loadSessions() {
    try {
        const stored = localStorage.getItem('guardrail_chat_sessions');
        if (stored) {
            const parsed = JSON.parse(stored);
            sessions = parsed.sessions || {};
            activeSessionId = parsed.activeSessionId || null;
        }
    } catch (e) {
        console.error('Error loading sessions:', e);
    }

    // If no sessions, create a default one
    const keys = Object.keys(sessions);
    if (keys.length === 0) {
        createNewSession();
    } else {
        if (!activeSessionId || !sessions[activeSessionId]) {
            activeSessionId = keys[0];
        }
        renderSidebar();
        renderActiveSession();
    }
}

function saveSessions() {
    try {
        localStorage.setItem('guardrail_chat_sessions', JSON.stringify({
            sessions,
            activeSessionId
        }));
    } catch (e) {
        console.error('Error saving sessions:', e);
    }
}

function createNewSession() {
    const sessionId = 'session_' + Date.now() + '_' + Math.random().toString(36).substr(2, 5);
    sessions[sessionId] = {
        id: sessionId,
        title: 'New Chat',
        messages: [],
        timestamp: Date.now()
    };
    activeSessionId = sessionId;
    saveSessions();
    renderSidebar();
    renderActiveSession();
}

function clearChat() {
    // Linked to the "New Chat" button
    createNewSession();
}

function renderSidebar() {
    const list = document.getElementById('historyList');
    if (!list) return;
    list.innerHTML = '';

    // Sort sessions by timestamp descending (newest first)
    const sortedSessions = Object.values(sessions).sort((a, b) => b.timestamp - a.timestamp);

    sortedSessions.forEach(session => {
        const item = document.createElement('div');
        item.className = 'history-item' + (session.id === activeSessionId ? ' active' : '');
        
        const titleSpan = document.createElement('span');
        titleSpan.className = 'history-title';
        titleSpan.innerText = '💬 ' + session.title;
        titleSpan.onclick = () => switchSession(session.id);

        const deleteBtn = document.createElement('button');
        deleteBtn.className = 'btn-delete-chat';
        deleteBtn.innerText = '❌';
        deleteBtn.onclick = (e) => {
            e.stopPropagation();
            deleteSession(session.id);
        };

        item.appendChild(titleSpan);
        item.appendChild(deleteBtn);
        list.appendChild(item);
    });
}

function switchSession(sessionId) {
    if (sessionId === activeSessionId) return;
    activeSessionId = sessionId;
    saveSessions();
    renderSidebar();
    renderActiveSession();
}

function deleteSession(sessionId) {
    if (sessions[sessionId]) {
        delete sessions[sessionId];
        const keys = Object.keys(sessions);
        if (keys.length === 0) {
            createNewSession();
            return;
        }
        if (activeSessionId === sessionId) {
            // Switch to the newest remaining session
            const sorted = Object.values(sessions).sort((a, b) => b.timestamp - a.timestamp);
            activeSessionId = sorted[0].id;
        }
        saveSessions();
        renderSidebar();
        renderActiveSession();
    }
}

function renderActiveSession() {
    const thread = document.getElementById('chatThread');
    const welcome = document.getElementById('welcomeContainer');
    
    thread.innerHTML = '';
    
    const activeSession = sessions[activeSessionId];
    if (!activeSession || activeSession.messages.length === 0) {
        thread.style.display = 'none';
        welcome.style.display = 'block';
    } else {
        welcome.style.display = 'none';
        thread.style.display = 'flex';
        
        activeSession.messages.forEach(msg => {
            appendMessageToUI(msg.content, msg.role, msg.id, msg.details);
        });
    }
}

async function sendMessage() {
    const input = document.getElementById('userInput');
    const text = input.value.trim();
    if (!text) return;

    input.value = '';
    input.style.height = 'auto';
    input.style.overflowY = 'hidden';

    const currentSession = sessions[activeSessionId];
    if (!currentSession) return;

    // Toggle Send button to Stop button
    const sendBtn = document.getElementById('sendBtn');
    if (sendBtn) {
        sendBtn.innerHTML = '<svg viewBox="0 0 24 24" width="18" height="18" fill="currentColor"><rect x="4" y="4" width="16" height="16" rx="2" /></svg>';
        sendBtn.setAttribute('onclick', 'stopGeneration()');
        sendBtn.title = 'Stop Generation';
    }

    activeAbortController = new AbortController();
    const { signal } = activeAbortController;

    // Hide welcome block and show thread
    document.getElementById('welcomeContainer').style.display = 'none';
    const thread = document.getElementById('chatThread');
    thread.style.display = 'flex';

    // If it's a new session, update title based on first query
    if (currentSession.title === 'New Chat') {
        currentSession.title = text.length > 22 ? text.substring(0, 19) + '...' : text;
    }
    
    // Update timestamp to float this session to top
    currentSession.timestamp = Date.now();

    // Generate random message IDs
    const userMsgId = 'msg_user_' + Date.now();
    const assistantMsgId = 'msg_ast_' + Date.now();

    // Save user message to state
    currentSession.messages.push({
        id: userMsgId,
        role: 'user',
        content: text
    });
    saveSessions();
    renderSidebar();

    // Append User message to UI
    appendMessageToUI(text, 'user', userMsgId);

    // Append loading placeholder
    const loaderId = appendMessageToUI('Thinking...', 'assistant loading', 'loader_' + Date.now());

    // Build message history for the AI memory
    const messageHistory = currentSession.messages
        .filter(msg => msg.role === 'user' || msg.role === 'assistant')
        .map(msg => ({
            role: msg.role,
            content: msg.content
        }));

    try {
        const response = await fetch('/v1/shield/generate', {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify({ 
                prompt: text, 
                messages: messageHistory,
                stream: true 
            }),
            signal: signal
        });
        
        if (!response.ok) {
            throw new Error(`HTTP error! status: ${response.status}`);
        }

        const reader = response.body.getReader();
        const decoder = new TextDecoder();
        let buffer = '';
        let hasAppendedAssistantMessage = false;
        let accumulatedText = '';

        // Remove loading placeholder
        const loaderEl = document.getElementById(loaderId);
        if (loaderEl) loaderEl.remove();

        while (true) {
            const { value, done } = await reader.read();
            if (done) break;

            buffer += decoder.decode(value, { stream: true });
            const lines = buffer.split('\n');
            buffer = lines.pop(); // Keep partial line in buffer

            for (const line of lines) {
                const trimmed = line.trim();
                if (trimmed.startsWith('data: ')) {
                    const jsonStr = trimmed.slice(6);
                    try {
                        const data = JSON.parse(jsonStr);
                        if (data.safe) {
                            if (data.response) {
                                accumulatedText += data.response;
                                
                                if (!hasAppendedAssistantMessage) {
                                    // Append empty assistant message to state and UI
                                    currentSession.messages.push({
                                        id: assistantMsgId,
                                        role: 'assistant',
                                        content: accumulatedText
                                    });
                                    appendMessageToUI(accumulatedText, 'assistant', assistantMsgId);
                                    hasAppendedAssistantMessage = true;
                                } else {
                                    // Update state
                                    currentSession.messages[currentSession.messages.length - 1].content = accumulatedText;
                                    
                                    // Update UI
                                    const msgElement = document.getElementById(assistantMsgId);
                                    if (msgElement) {
                                        const bubble = msgElement.querySelector('.bubble');
                                        bubble.innerHTML = parseMarkdown(accumulatedText);
                                        highlightCode(bubble);
                                    }
                                }
                                
                                // Auto-scroll
                                const viewport = document.querySelector('.chat-viewport');
                                viewport.scrollTop = viewport.scrollHeight;
                            }
                        } else {
                            // Blocked mid-stream or initially
                            if (hasAppendedAssistantMessage) {
                                // Remove partial message from state
                                currentSession.messages.pop();
                                const msgElement = document.getElementById(assistantMsgId);
                                if (msgElement) msgElement.remove();
                            }
                            
                            const status = (data.details && data.details.gateway_status) ? data.details.gateway_status : 'validation_failed';
                            const blockText = `⚠️ Request Blocked: Response violated guardrail safety parameters (${status})`;
                            
                            // Save blocked message to state
                            currentSession.messages.push({
                                id: assistantMsgId,
                                role: 'blocked',
                                content: blockText,
                                details: data.details
                            });
                            
                            appendMessageToUI(blockText, 'blocked', assistantMsgId, data.details);
                            saveSessions();
                            return;
                        }
                    } catch (err) {
                        console.error('Error parsing SSE JSON:', err);
                    }
                }
            }
        }
        
        // Final save on successful completion
        saveSessions();
    } catch (e) {
        const loaderEl = document.getElementById(loaderId);
        if (loaderEl) loaderEl.remove();
        
        if (e.name === 'AbortError') {
            if (accumulatedText) {
                // Remove loading styling from bubble
                const msgElement = document.getElementById(assistantMsgId);
                if (msgElement) {
                    msgElement.classList.remove('loading');
                }
            } else {
                const stopText = 'Generation stopped.';
                currentSession.messages.push({
                    id: 'stop_' + Date.now(),
                    role: 'assistant',
                    content: stopText
                });
                appendMessageToUI(stopText, 'assistant', 'stop_' + Date.now());
            }
            saveSessions();
        } else {
            const errorText = '❌ Error: Failed to communicate with the guardrail gateway.';
            currentSession.messages.push({
                id: 'err_' + Date.now(),
                role: 'blocked',
                content: errorText
            });
            appendMessageToUI(errorText, 'blocked', 'err_' + Date.now());
            saveSessions();
        }
    } finally {
        // Restore Send button UI
        const sendBtn = document.getElementById('sendBtn');
        if (sendBtn) {
            sendBtn.innerHTML = '<svg viewBox="0 0 24 24" width="18" height="18" fill="currentColor"><path d="M2.01 21L23 12 2.01 3 2 10l15 2-15 2z" /></svg>';
            sendBtn.setAttribute('onclick', 'sendMessage()');
            sendBtn.title = 'Send Message';
        }
        activeAbortController = null;
    }
}

function stopGeneration() {
    if (activeAbortController) {
        activeAbortController.abort();
        activeAbortController = null;
    }
}

function appendMessageToUI(text, type, msgId, details = null) {
    const thread = document.getElementById('chatThread');
    if (!thread) return msgId;

    const msgDiv = document.createElement('div');
    msgDiv.className = `message ${type}`;
    msgDiv.id = msgId;

    let contentHtml = '';
    if (type === 'user') {
        contentHtml = escapeHTML(text);
    } else if (type === 'assistant') {
        contentHtml = parseMarkdown(text);
    } else if (type === 'blocked') {
        contentHtml = escapeHTML(text);
        if (details && details.input_guard && details.input_guard.matched_rule !== 'none') {
            contentHtml += `<div style="font-size: 11px; margin-top: 5px; opacity: 0.8;">Matched rule: <strong>${escapeHTML(details.input_guard.matched_rule)}</strong></div>`;
        }
    } else {
        contentHtml = escapeHTML(text);
    }

    msgDiv.innerHTML = `<div class="bubble">${contentHtml}</div>`;
    thread.appendChild(msgDiv);

    if (type === 'assistant') {
        highlightCode(msgDiv);
    }

    // Auto-scroll
    const viewport = document.querySelector('.chat-viewport');
    viewport.scrollTop = viewport.scrollHeight;
    
    return msgId;
}

function parseMarkdown(text) {
    if (typeof marked !== 'undefined' && typeof marked.parse === 'function') {
        return marked.parse(text);
    }
    return escapeHTML(text).replace(/\n/g, '<br>');
}

function highlightCode(element) {
    if (typeof hljs !== 'undefined' && typeof hljs.highlightElement === 'function') {
        element.querySelectorAll('pre code').forEach((block) => {
            hljs.highlightElement(block);
        });
    }
}

function escapeHTML(str) {
    return str
        .replace(/&/g, '&amp;')
        .replace(/</g, '&lt;')
        .replace(/>/g, '&gt;')
        .replace(/"/g, '&quot;')
        .replace(/'/g, '&#039;');
}