{% 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;
    }

    .chat-output {
        min-height: 150px;
        border: 1px solid #ddd;
        padding: 10px;
        margin-bottom: 10px;
        overflow-y: auto;
    }

    .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;
    }
</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>
                <p><strong>Generated:</strong> <span id="{{ split }}-generated"></span></p>
                <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 %}

        <div id="load-chat-container">
            <button id="load-chat-button">Load Chat with {{ model_name }}</button>
        </div>
        <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="">
                <!-- Removed the Update System Message button -->
            </div>

            <div id="chat-output" class="chat-output"></div>
            <input type="text" id="chat-input" class="chat-input" placeholder="Enter your message..." disabled>
        </div>
    </div>
</div>

<script>
    var generatedData = {{ generated_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;
    }

    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);

                fetch('/chat', {
                    method: 'POST',
                    headers: {
                        'Content-Type': 'application/x-www-form-urlencoded',
                    },
                    body: 'message=' + encodeURIComponent(message)
                })
                    .then(response => response.json())
                    .then(data => {
                        addMessageToChat('{{ model_name }}: ' + 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;
                }
            });
    }

    // 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';
                    document.getElementById('chat-input').disabled = false;

                    console.log("Chat interface loaded successfully.");
                    initializeChat();
                }
            })
            .catch(error => {
                console.error("Error during /load_model fetch:", error);
            });
    });


    function addMessageToChat(message) {
        var chatOutput = document.getElementById('chat-output');
        var newMessage = document.createElement('p');
        newMessage.textContent = message;
        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();
        }
    });

</script>

{% endblock %}