{% extends "base.html" %}

{% block title %}Training Run Visualization: {{ run }}{% endblock %}

{% block head %}
<style>
    table {
        width: 80%;
        margin: 20px auto;
        border-collapse: collapse;
    }

    th,
    td {
        border: 1px solid #ddd;
        padding: 8px;
        text-align: left;
    }

    th {
        background-color: #f2f2f2;
    }

    .group-header {
        background-color: #e0e0e0;
        font-weight: bold;
    }

    .container {
        display: flex;
        gap: 20px;
        align-items: flex-start;
        justify-content: space-between;
        margin-right: 20px;
    }

    .column {
        flex: 1;
        max-width: 80vw;
        padding: 20px;
        background-color: #f8f8f8;
        border: 1px solid #ddd;
        border-radius: 10px;
        box-shadow: 2px 2px 5px rgba(0, 0, 0, 0.1);
    }

    .column:first-child {
        flex: 0 0 30%;
    }

    .column:last-child {
        flex: 0 0 65%;
    }

    .generated-text-container {
        margin-bottom: 20px;
    }

    .generated-text-controls {
        margin-top: 10px;
        display: flex;
        align-items: center;
        gap: 10px;
    }

    .generated-text-controls button {
        margin: 0;
        padding: 8px 16px;
        background-color: #4CAF50;
        /* Green */
        color: white;
        border: none;
        border-radius: 5px;
        cursor: pointer;
        transition: background-color 0.3s ease;
    }

    .generated-text-controls button:hover {
        background-color: #367c39;
        /* Darker green */
    }

    .generated-text-controls input {
        width: 80px;
        padding: 8px;
        margin: 0;
        border: 1px solid #ddd;
        border-radius: 5px;
        text-align: center;
    }

    .generated-text-display {
        min-height: 150px;
    }

    .generated-text-display p span {
        display: block;
        overflow-y: auto;
        max-height: 100px;
        /* Adjust as needed */
        border: 1px solid #ddd;
        padding: 5px;
        margin-bottom: 5px;
        min-height: 100px;
        width: 95%;
        max-width: 95%;
        overflow-x: auto;
        /* Add horizontal scroll for long text */
        white-space: pre-wrap;
        /* Preserve line breaks and wrap text */
        word-break: break-all;
        /* Force break words even without spaces */
    }

    .index-display {
        font-size: 1.1em;
        color: #555;
    }

    button {
        margin: 0;
        padding: 8px 16px;
        background-color: #4CAF50;
        /* Green */
        color: white;
        border: none;
        border-radius: 5px;
        cursor: pointer;
        transition: background-color 0.3s ease;
    }

    button:hover {
        background-color: #367c39;
        /* Darker green */
    }

    .top-controls {
        display: flex;
        justify-content: space-between;
        align-items: center;
        margin-bottom: 20px;
    }

    .top-controls>div {
        display: flex;
        flex-direction: column;
        align-items: flex-start;
    }

    .top-controls h1 {
        margin-bottom: 5px;
    }

    .top-controls button {
        margin-right: 20px;
        padding: 12px 24px;
        font-size: 1.2em;
    }

    .performance-table {
        width: auto;
        /* Adjust width to fit content */
        margin: 20px 0;
        border-collapse: collapse;
    }

    .performance-table th,
    .performance-table td {
        border: 1px solid #ddd;
        padding: 8px;
        text-align: left;
        font-size: 0.9em;
        /* Reduce font size */
    }

    .performance-table th {
        background-color: #f2f2f2;
    }

    .chat-container {
        margin-bottom: 20px;
        border: 1px solid #ddd;
        border-radius: 10px;
        padding: 15px;
        background-color: #fafafa;
        box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
    }

    .chat-container h2 {
        margin-top: 0;
        margin-bottom: 15px;
        color: #333;
        border-bottom: 1px solid #eee;
        padding-bottom: 10px;
    }

    .chat-output {
        min-height: 200px;
        max-height: 400px;
        border: 1px solid #ddd;
        border-radius: 8px;
        padding: 10px;
        margin-bottom: 10px;
        overflow-y: auto;
        background-color: white;
    }

    .chat-output p {
        margin: 8px 0;
        white-space: pre-wrap;
        word-wrap: break-word;
        line-height: 1.4;
    }

    .chat-message {
        padding: 8px 12px;
        border-radius: 10px;
        margin-bottom: 10px;
        max-width: 85%;
    }

    .chat-message.user {
        background-color: #e3f2fd;
        margin-left: auto;
    }

    .chat-message.assistant {
        background-color: #f5f5f5;
        margin-right: auto;
    }

    .chat-message.system {
        background-color: #fffde7;
        font-style: italic;
        width: 100%;
        text-align: center;
        font-size: 0.9em;
    }

    .chat-message .message-content {
        white-space: pre-wrap !important;
        word-wrap: break-word;
        overflow-wrap: break-word;
    }

    .chat-input {
        width: 100%;
        padding: 10px;
        border: 1px solid #ddd;
        border-radius: 5px;
    }

    .system-msg-container {
        margin-bottom: 10px;
    }

    .system-msg-container label {
        display: block;
        margin-bottom: 5px;
    }

    .system-msg-container input {
        width: 100%;
        padding: 8px;
        border: 1px solid #ddd;
        border-radius: 5px;
    }

    /* Removed system-msg-container button styles */

    #eval-folder-selection {
        margin-bottom: 20px;
    }

    #eval-folder-selection label {
        display: block;
        margin-bottom: 5px;
        font-weight: bold;
    }

    #eval-folder-selection select {
        width: 300px;
        padding: 8px;
        border: 1px solid #ddd;
        border-radius: 5px;
    }

    .hypernetwork-container {
        margin-top: 20px;
        padding: 15px;
        background-color: #f0f8ff;
        border: 1px solid #add8e6;
        border-radius: 5px;
    }

    .hypernetwork-container h3 {
        margin-top: 0;
        margin-bottom: 15px;
        color: #0066cc;
    }

    .hypernetwork-container label {
        display: block;
        margin-bottom: 5px;
        font-weight: bold;
    }

    .hypernetwork-container select,
    .hypernetwork-container textarea {
        width: 100%;
        padding: 8px;
        margin-bottom: 15px;
        border: 1px solid #add8e6;
        border-radius: 5px;
    }

    .hypernetwork-container textarea {
        height: 120px;
        resize: vertical;
    }

    .chat-controls {
        margin-top: 10px;
        display: flex;
        gap: 10px;
    }

    .secondary-button {
        background-color: #607d8b;
        color: white;
        padding: 8px 16px;
        border: none;
        border-radius: 5px;
        cursor: pointer;
        transition: background-color 0.3s;
    }

    .secondary-button:hover {
        background-color: #455a64;
    }

    .instruction-text {
        margin: 0 0 15px 0;
        font-style: italic;
        color: #555;
        font-size: 0.9em;
    }

    .ui-sequence {
        margin-bottom: 10px;
        padding: 10px;
        background-color: #f5f5f5;
        border-left: 4px solid #4CAF50;
        color: #333;
        font-size: 0.9em;
    }

    .context-field {
        margin-bottom: 15px;
        padding: 15px;
        border: 1px solid #e0e0e0;
        border-radius: 5px;
        background-color: #f9f9f9;
    }

    .context-field label {
        margin-bottom: 8px;
        color: #333;
        font-weight: bold;
    }

    .context-input {
        width: 100%;
        height: 100px;
        resize: vertical;
        padding: 8px;
        border: 1px solid #ccc;
        border-radius: 5px;
    }

    .context-controls {
        display: flex;
        gap: 10px;
        margin-bottom: 15px;
    }

    .tertiary-button {
        background-color: #ff5252;
        color: white;
        padding: 8px 16px;
        border: none;
        border-radius: 5px;
        cursor: pointer;
        transition: background-color 0.3s;
    }

    .tertiary-button:hover {
        background-color: #d32f2f;
    }

    .tertiary-button:disabled {
        background-color: #ffcdd2;
        cursor: not-allowed;
    }

    .context-header {
        margin-bottom: 15px;
    }

    .context-header h4 {
        margin: 0 0 5px 0;
        color: #333;
    }

    .chat-input.textarea {
        min-height: 60px;
        max-height: 200px;
        resize: vertical;
        overflow-y: auto;
        font-family: inherit;
        font-size: inherit;
        padding: 10px;
        box-sizing: border-box;
    }

    .chat-input-container {
        display: flex;
        gap: 10px;
        margin-top: 10px;
    }

    .chat-input-container textarea {
        flex: 1;
    }

    .primary-button {
        background-color: #2196F3;
        color: white;
        padding: 8px 16px;
        border: none;
        border-radius: 5px;
        cursor: pointer;
        transition: background-color 0.3s;
        align-self: flex-end;
    }

    .primary-button:hover {
        background-color: #0d8aee;
    }

    .primary-button:disabled {
        background-color: #bbdefb;
        cursor: not-allowed;
    }

    .model-comparison {
        display: flex;
        flex-wrap: wrap;
        gap: 20px;
        margin-bottom: 15px;
    }

    .model-output {
        flex: 1;
        min-width: 300px;
        border: 1px solid #e0e0e0;
        border-radius: 5px;
        padding: 10px;
        background-color: #f9f9f9;
    }

    .model-output p {
        margin: 0;
    }

    .model-output span {
        display: block;
        padding: 10px;
        background-color: white;
        border: 1px solid #eee;
        border-radius: 4px;
        min-height: 80px;
        max-height: 150px;
        overflow-y: auto;
        white-space: pre-wrap;
        margin-top: 5px;
    }

    .highlight-diff {
        background-color: #fff8e1;
    }
</style>
{% endblock %}

{% block content %}
<div class="top-controls">
    <div>
        <h1>Training Run: {{ run }}</h1>
        <p>Command: <span id="command"></span></p>
    </div>
    <button onclick="goBack()">Go Back</button>
</div>

<div class="container">
    <div class="column">
        {% if eval_folders %} <!-- Conditionally show folder selection if eval_folders exist -->
        <div id="eval-folder-selection">
            <label for="eval_folder">Select Evaluation Folder:</label>
            <select id="eval_folder" onchange="loadEvalFolder()">
                <option value="">-- Select Folder --</option>
                {% for folder in eval_folders %}
                <option value="{{ folder }}" {% if selected_eval_folder==folder %}selected{% endif %}>{{ folder }}
                </option>
                {% endfor %}
            </select>
        </div>
        {% endif %}

        {% if data %}
        {% for group, results in data.items() %}
        <h2>{{ group }} ({% if selected_eval_folder and selected_eval_folder != "root" %}{{ selected_eval_folder }}{%
            else %}root{% endif %})</h2>
        <table class="performance-table">
            <thead>
                <tr>
                    <th>Metric</th>
                    <th>Value</th>
                </tr>
            </thead>
            <tbody>
                {% for metric, value in results.items() %}
                <tr>
                    <td>{{ metric }}</td>
                    {% if value is mapping %}
                    <td>
                        <table class="performance-table">
                            <thead>
                                <tr>
                                    <th>Sub-Metric</th>
                                    <th>Value</th>
                                </tr>
                            </thead>
                            <tbody>
                                {% for submetric, subvalue in value.items() %}
                                <tr>
                                    <td>{{ submetric }}</td>
                                    <td>{{ subvalue }}</td>
                                </tr>
                                {% endfor %}
                            </tbody>
                        </table>
                    </td>
                    {% else %}
                    <td>{{ value }}</td>
                    {% endif %}
                </tr>
                {% endfor %}
            </tbody>
        </table>
        {% endfor %}
        {% endif %}
    </div>

    <div class="column">
        {% if generated_data %}
        {% for split, data in generated_data.items() %}
        {% if data %}
        <div class="generated-text-container">
            <h2>Generated Text ({{ split }} - {% if selected_eval_folder and selected_eval_folder != "root" %}{{
                selected_eval_folder }}{% else %}root{% endif %})</h2>
            <div id="{{ split }}-text" class="generated-text-display">
                <p id="{{ split }}-context-container" style="display: none;">
                    <strong>Context:</strong> <span id="{{ split }}-context"></span>
                </p>
                <p><strong>Input:</strong> <span id="{{ split }}-input"></span></p>

                <!-- Display both models with enhanced styling -->
                <div class="model-comparison">
                    <div class="model-output">
                        <p><strong>Modulated Model:</strong> <span id="{{ split }}-generated"></span></p>
                    </div>
                    <div class="model-output">
                        <p><strong>Base Model ({{ model_name }}):</strong> <span id="{{ split }}-base-generated">Not
                                available</span></p>
                    </div>
                    <div class="model-output">
                        <p><strong>Base Model (No Context - {{ model_name }}):</strong> <span
                                id="{{ split }}-base-no-context-generated">Not available</span></p>
                    </div>
                </div>

                <p><strong>Label:</strong> <span id="{{ split }}-label"></span></p>
            </div>
            <div class="generated-text-controls">
                <button onclick="prev('{{ split }}')">Prev</button>
                <input type="number" id="{{ split }}-index" value="1" min="1" max="{{ data|length }}"
                    onchange="updateText('{{ split }}')">
                <span class="index-display">/ {{ data|length }}</span>
                <button onclick="next('{{ split }}')">Next</button>
                <button onclick="randomize('{{ split }}')">Random</button>
            </div>
        </div>
        {% endif %}
        {% endfor %}
        {% elif selected_eval_folder %}
        <p>No generated text data found in {% if selected_eval_folder and selected_eval_folder != "root" %}{{
            selected_eval_folder }}{% else %}root{% endif %}.</p>
        {% endif %}

        <!-- Step-by-step guide -->
        <div class="ui-sequence">
            <h3>How to use this interface:</h3>
            <ol>
                <li>First, click "Load Chat with {{ model_name }}" to initialize the base model</li>
                <li>You can then chat directly with the base model</li>
                <li>Alternatively, you can select a hypernetwork checkpoint to use context-informed responses</li>
                <li>To reset the conversation at any time, click the "Reset Chat" button</li>
            </ol>
        </div>

        <div id="load-chat-container">
            <button id="load-chat-button">Load Chat with {{ model_name }}</button>
        </div>

        <!-- Hypernetwork Section - Initially hidden -->
        {% if checkpoints %}
        <div class="hypernetwork-container" id="hypernetwork-container" style="display: none;">
            <h3>Hypernetwork Context Integration</h3>
            <p class="instruction-text">Apply a hypernetwork to modulate the base model with context information</p>
            <label for="checkpoint-select">Select Checkpoint:</label>
            <select id="checkpoint-select">
                <option value="">-- Select Checkpoint --</option>
                {% for checkpoint in checkpoints %}
                <option value="{{ checkpoint }}">{{ checkpoint }}</option>
                {% endfor %}
            </select>

            <div id="contexts-container">
                <div class="context-header">
                    <h4>Context Inputs</h4>
                    <p class="instruction-text">Add one or more context inputs to generate corresponding LoRAs.</p>
                </div>

                <div class="context-controls">
                    <button id="add-context-button" class="secondary-button">Add Context</button>
                    <button id="remove-context-button" class="tertiary-button" disabled>Remove Last</button>
                </div>

                <div class="context-fields">
                    <div class="context-field">
                        <label for="context-input-0">Context 1:</label>
                        <textarea id="context-input-0" class="context-input"
                            placeholder="Enter context information here. This will be used by the hypernetwork to generate a LoRA for the base model."></textarea>
                    </div>
                </div>
            </div>

            <button id="apply-hypernetwork-button">Apply Hypernetwork</button>
            <div id="hypernetwork-status"></div>
        </div>
        {% endif %}

        <div class="chat-container" id="chat-container" style="display: none;">
            <h2>Chat with <span id="chat-model-name"></span></h2>

            <div class="system-msg-container">
                <label for="system-msg">System Message:</label>
                <input type="text" id="system-msg" class="chat-input" placeholder="Enter system message..." value="">
                <div class="chat-controls">
                    <button id="reset-chat-button" class="secondary-button">Reset Chat</button>
                </div>
            </div>

            <div id="chat-output" class="chat-output"></div>

            <div class="chat-input-container">
                <textarea id="chat-input" class="chat-input textarea"
                    placeholder="Enter your message... (Shift+Enter for new line)" disabled></textarea>
                <button id="send-message-button" class="primary-button" disabled>Send</button>
            </div>
        </div>
    </div>
</div>

<script>
    var generatedData = {{ generated_data | tojson }};
    var baseModelData = {{ base_model_data | tojson }};
    var baseModelNoContextData = {{ base_model_no_context_data | tojson }};

    function updateText(split) {
        var index = parseInt(document.getElementById(split + '-index').value) - 1;
        var data = generatedData[split][index];

        // Check if 'context' exists in the data and update it
        if ('context' in data) {
            document.getElementById(split + '-context').textContent = data.context;
            document.getElementById(split + '-context-container').style.display = 'block';
        } else {
            document.getElementById(split + '-context-container').style.display = 'none';
        }

        document.getElementById(split + '-input').textContent = data.input;
        document.getElementById(split + '-generated').textContent = data.generated;
        document.getElementById(split + '-label').textContent = data.label;

        // Update base model output if available
        var baseModelOutput = document.getElementById(split + '-base-generated');
        if (baseModelData && baseModelData[split] && baseModelData[split].length > index) {
            baseModelOutput.textContent = baseModelData[split][index].generated;
        } else {
            baseModelOutput.textContent = "Not available";
        }

        // Update base model (no context) output if available
        var baseModelNoContextOutput = document.getElementById(split + '-base-no-context-generated');
        if (baseModelNoContextData && baseModelNoContextData[split] && baseModelNoContextData[split].length > index) {
            baseModelNoContextOutput.textContent = baseModelNoContextData[split][index].generated;
        } else {
            baseModelNoContextOutput.textContent = "Not available";
        }
    }

    function next(split) {
        var indexInput = document.getElementById(split + '-index');
        var index = parseInt(indexInput.value);
        if (index < generatedData[split].length) {
            indexInput.value = index + 1;
            updateText(split);
        }
    }

    function prev(split) {
        var indexInput = document.getElementById(split + '-index');
        var index = parseInt(indexInput.value);
        if (index > 1) {
            indexInput.value = index - 1;
            updateText(split);
        }
    }

    function randomize(split) {
        var randomIndex = Math.floor(Math.random() * generatedData[split].length) + 1;
        document.getElementById(split + '-index').value = randomIndex;
        updateText(split);
    }

    // Initialize text for each split using JavaScript if generatedData is available
    window.addEventListener('DOMContentLoaded', (event) => {
        if (generatedData && Object.keys(generatedData).length > 0) {
            for (const split in generatedData) {
                if (generatedData[split].length > 0) {
                    updateText(split);
                }
            }
        }
    });


    function goBack() {
        window.history.back();
    }

    function getCommand(run) {
        fetch('/get_command/' + run)
            .then(response => response.json())
            .then(data => {
                if (data.command) {
                    document.getElementById('command').textContent = data.command;
                } else {
                    document.getElementById('command').textContent = 'Command not found.';
                }
            })
            .catch(error => {
                console.error("Error fetching command:", error);
                document.getElementById('command').textContent = 'Error loading command.';
            });
    }

    // Call getCommand when the page loads
    window.addEventListener('DOMContentLoaded', (event) => {
        getCommand('{{ run }}');
    });

    function sendMessage() {
        const systemMsg = document.getElementById('system-msg').value;
        fetch('/update_system_msg', {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json',
            },
            body: JSON.stringify({ system_msg: systemMsg })
        })
            .then(response => response.json())
            .then(data => {
                if (data.success) {
                    console.log('System message updated successfully.');
                } else {
                    console.log('No update needed for system message.');
                }
                // Proceed to send the chat message
                const message = document.getElementById('chat-input').value;
                document.getElementById('chat-input').value = '';
                addMessageToChat('You: ' + message);

                // Prepare form data
                let formData = new FormData();
                formData.append('message', message);

                // Check if we're using hypernetwork and include contexts
                if (document.getElementById('chat-model-name').textContent.includes('Hypernetwork')) {
                    const contextInputs = document.querySelectorAll('.context-input');

                    // Add all contexts to the form data
                    Array.from(contextInputs).forEach(input => {
                        formData.append('contexts[]', input.value);
                    });
                }

                fetch('/chat', {
                    method: 'POST',
                    body: formData
                })
                    .then(response => response.json())
                    .then(data => {
                        const modelName = document.getElementById('chat-model-name').textContent || '{{ model_name }}';
                        addMessageToChat(modelName + ': ' + data.response);
                    })
                    .catch(error => {
                        console.error("Error during chat fetch:", error);
                    });
            })
            .catch(error => {
                console.error('Error updating system message:', error);
            });
    }

    function initializeChat() {
        fetch('/get_system_msg')
            .then(response => response.json())
            .then(data => {
                if (data.system_msg) {
                    document.getElementById('system-msg').value = data.system_msg;
                    chat_history[0]["content"] = data.system_msg;
                }
            });
    }

    // Handle hypernetwork checkpoint loading
    document.addEventListener('DOMContentLoaded', function () {
        const applyHypernetworkButton = document.getElementById('apply-hypernetwork-button');
        if (applyHypernetworkButton) {
            applyHypernetworkButton.addEventListener('click', function () {
                const checkpoint = document.getElementById('checkpoint-select').value;
                const contextInputs = document.querySelectorAll('.context-input');

                // Collect all contexts
                const contexts = Array.from(contextInputs).map(input => input.value);

                if (!checkpoint) {
                    alert('Please select a checkpoint first.');
                    return;
                }

                // Check if all contexts are empty
                const allEmpty = contexts.every(context => !context.trim());
                if (allEmpty) {
                    if (!confirm('All context fields are empty. Are you sure you want to continue without any context?')) {
                        return;
                    }
                }

                const statusEl = document.getElementById('hypernetwork-status');
                statusEl.textContent = 'Loading checkpoint...';
                statusEl.style.color = '#1565c0'; // Blue to indicate in progress

                // Disable the button while loading
                applyHypernetworkButton.disabled = true;
                applyHypernetworkButton.textContent = 'Loading...';

                // Build the form data with all contexts
                const formData = new FormData();
                formData.append('run', '{{ run }}');
                formData.append('checkpoint', checkpoint);

                // Add all contexts to the form data
                contexts.forEach(context => {
                    formData.append('contexts[]', context);
                });

                fetch('/load_checkpoint', {
                    method: 'POST',
                    body: formData
                })
                    .then(response => response.json())
                    .then(data => {
                        // Re-enable the button
                        applyHypernetworkButton.disabled = false;
                        applyHypernetworkButton.textContent = 'Apply Hypernetwork';

                        if (data.success) {
                            const contextCount = data.contexts_processed || 0;
                            statusEl.textContent = `Success: ${data.message} (${contextCount} contexts processed)`;
                            statusEl.style.color = 'green';

                            // Show the chat container if it's hidden
                            document.getElementById('chat-container').style.display = 'block';

                            // Enable the chat input
                            const chatInput = document.getElementById('chat-input');
                            const sendButton = document.getElementById('send-message-button');
                            if (chatInput) {
                                chatInput.disabled = false;
                                // Still disabled until text is entered
                                if (sendButton) sendButton.disabled = true;
                            }

                            document.getElementById('chat-model-name').textContent =
                                'Hypernetwork: ' + checkpoint;

                            // Clear any existing chat and add an instruction
                            document.getElementById('chat-output').innerHTML = '';
                            addMessageToChat('System: Hypernetwork applied successfully. Your messages will now be processed through the context-aware model.');
                            if (contextCount > 1) {
                                addMessageToChat(`System: ${contextCount} contexts were processed to generate combined LoRAs.`);
                            }

                        } else {
                            statusEl.textContent = 'Error: ' + data.error;
                            statusEl.style.color = 'red';
                        }
                    })
                    .catch(error => {
                        // Re-enable the button on error
                        applyHypernetworkButton.disabled = false;
                        applyHypernetworkButton.textContent = 'Apply Hypernetwork';

                        console.error('Error loading checkpoint:', error);
                        statusEl.textContent = 'Error: ' + error.message;
                        statusEl.style.color = 'red';
                    });
            });
        }
    });

    // Call initializeChat when chat is loaded
    document.getElementById('load-chat-button').addEventListener('click', function () {
        fetch('/load_model', {
            method: 'POST',
            headers: {
                'Content-Type': 'application/x-www-form-urlencoded',
            },
            body: 'run=' + encodeURIComponent('{{ run }}')
        })
            .then(response => {
                console.log("Received response status:", response.status);
                if (!response.ok) {
                    throw new Error('Network response was not ok');
                }
                return response.json();
            })
            .then(data => {
                console.log("Parsed data:", data);
                if (data.error) {
                    alert(data.error);
                } else {
                    document.getElementById('chat-model-name').textContent = data.model_name;
                    document.getElementById('load-chat-container').style.display = 'none';
                    document.getElementById('chat-container').style.display = 'block';

                    // Enable the input and send button
                    const chatInput = document.getElementById('chat-input');
                    const sendButton = document.getElementById('send-message-button');
                    chatInput.disabled = false;
                    if (sendButton) sendButton.disabled = true; // Still disabled until text is entered

                    // Show the hypernetwork container now that base model is loaded
                    const hypernetworkContainer = document.getElementById('hypernetwork-container');
                    if (hypernetworkContainer) {
                        hypernetworkContainer.style.display = 'block';
                    }

                    console.log("Chat interface loaded successfully.");
                    initializeChat();

                    // Add a message about UI flow
                    addMessageToChat('System: Base model loaded successfully. You can now chat directly or apply a hypernetwork with context information.');
                }
            })
            .catch(error => {
                console.error("Error during /load_model fetch:", error);
            });
    });


    function addMessageToChat(message) {
        var chatOutput = document.getElementById('chat-output');
        var newMessage = document.createElement('div');

        // Determine message type
        if (message.startsWith('You: ')) {
            newMessage.className = 'chat-message user';
            const messageContent = document.createElement('div');
            messageContent.className = 'message-content';
            messageContent.textContent = message.replace('You: ', '');
            newMessage.appendChild(messageContent);
        } else if (message.startsWith('System: ')) {
            newMessage.className = 'chat-message system';
            const messageContent = document.createElement('div');
            messageContent.className = 'message-content';
            messageContent.textContent = message.replace('System: ', '');
            newMessage.appendChild(messageContent);
        } else {
            newMessage.className = 'chat-message assistant';
            const colonIndex = message.indexOf(': ');
            if (colonIndex !== -1) {
                const sender = message.substring(0, colonIndex);
                const content = message.substring(colonIndex + 2);

                // Add a small label for the model name
                const modelLabel = document.createElement('div');
                modelLabel.className = 'model-label';
                modelLabel.style.fontSize = '0.8em';
                modelLabel.style.fontWeight = 'bold';
                modelLabel.style.marginBottom = '5px';
                modelLabel.style.color = '#555';
                modelLabel.textContent = sender;

                newMessage.appendChild(modelLabel);

                // Add the message content in a separate div
                const messageContent = document.createElement('div');
                messageContent.className = 'message-content';
                messageContent.textContent = content;
                newMessage.appendChild(messageContent);
            } else {
                const messageContent = document.createElement('div');
                messageContent.className = 'message-content';
                messageContent.textContent = message;
                newMessage.appendChild(messageContent);
            }
        }

        chatOutput.appendChild(newMessage);
        chatOutput.scrollTop = chatOutput.scrollHeight;
    }

    function loadEvalFolder() {
        const selectedFolder = document.getElementById('eval_folder').value;
        const runName = '{{ run }}';
        if (selectedFolder) {
            window.location.href = `/visualize/${runName}?eval_folder=${selectedFolder}`;
        }
    }


    // Bind the sendMessage function to the Enter key
    document.getElementById('chat-input').addEventListener('keypress', function (e) {
        if (e.key === 'Enter') {
            sendMessage();
        }
    });

    // Add the reset chat functionality
    document.addEventListener('DOMContentLoaded', function () {
        const resetChatButton = document.getElementById('reset-chat-button');
        if (resetChatButton) {
            resetChatButton.addEventListener('click', function () {
                // Get the current system message
                const systemMsg = document.getElementById('system-msg').value;

                // Clear the chat output display
                document.getElementById('chat-output').innerHTML = '';

                fetch('/reset_chat', {
                    method: 'POST',
                    headers: {
                        'Content-Type': 'application/x-www-form-urlencoded',
                    },
                    body: 'system_msg=' + encodeURIComponent(systemMsg)
                })
                    .then(response => response.json())
                    .then(data => {
                        if (data.success) {
                            console.log('Chat history reset successfully');
                            addMessageToChat('System: Chat history has been reset.');
                        } else {
                            console.error('Failed to reset chat history');
                        }
                    })
                    .catch(error => {
                        console.error('Error resetting chat history:', error);
                    });
            });
        }
    });

    // Add JavaScript to manage multiple context fields
    document.addEventListener('DOMContentLoaded', function () {
        const addContextButton = document.getElementById('add-context-button');
        const removeContextButton = document.getElementById('remove-context-button');
        const contextFieldsContainer = document.querySelector('.context-fields');

        let contextCount = 1; // We start with one context field

        // Add a new context field
        addContextButton.addEventListener('click', function () {
            const newFieldIndex = contextCount;
            contextCount++;

            const newField = document.createElement('div');
            newField.className = 'context-field';
            newField.innerHTML = `
                <label for="context-input-${newFieldIndex}">Context ${contextCount}:</label>
                <textarea id="context-input-${newFieldIndex}" class="context-input" 
                    placeholder="Enter additional context information here."></textarea>
            `;

            contextFieldsContainer.appendChild(newField);

            // Enable the remove button if we have more than one context
            if (contextCount > 1) {
                removeContextButton.disabled = false;
            }
        });

        // Remove the last context field
        removeContextButton.addEventListener('click', function () {
            if (contextCount > 1) {
                contextFieldsContainer.removeChild(contextFieldsContainer.lastChild);
                contextCount--;

                // Disable the remove button if we're back to just one context
                if (contextCount === 1) {
                    removeContextButton.disabled = true;
                }
            }
        });
    });

    // Add event handler for the send button
    document.addEventListener('DOMContentLoaded', function () {
        const sendButton = document.getElementById('send-message-button');
        const chatInput = document.getElementById('chat-input');

        if (sendButton && chatInput) {
            // Enable/disable the button based on chat input state
            const updateSendButtonState = function () {
                sendButton.disabled = chatInput.disabled || !chatInput.value.trim();
            };

            // Initial state
            updateSendButtonState();

            // Add event listener for input changes
            chatInput.addEventListener('input', updateSendButtonState);

            // Add event listener for send button click
            sendButton.addEventListener('click', function () {
                if (!chatInput.disabled && chatInput.value.trim()) {
                    sendMessage();
                }
            });

            // Update the event listener for chat input to handle Shift+Enter
            chatInput.addEventListener('keydown', function (e) {
                if (e.key === 'Enter') {
                    if (e.shiftKey) {
                        // Let Shift+Enter add a new line - no need to do anything
                        return;
                    } else {
                        // Just Enter triggers send
                        e.preventDefault();
                        if (!chatInput.disabled && chatInput.value.trim()) {
                            sendMessage();
                        }
                    }
                }
            });
        }
    });

</script>

{% endblock %}