import streamlit as st
import streamlit.components.v1 as components

st.title("Aider WebSocket Client")
st.write(
    "Interact with the Aider WebSocket server. Send commands, run files, and terminate running processes."
)

html_code = """
<div>
    <div id="messages" style="height: 300px; overflow-y: scroll; border: 1px solid #ccc; padding: 10px;"></div>
    <input id="repo_url" type="text" placeholder="Repository URL (optional)" style="margin-top: 10px; width: 300px;">
    <button onclick="connectToAider()" style="margin-left: 10px;">Connect to Aider</button>
    <input id="aider_input" type="text" placeholder="Send to Aider" style="margin-top: 10px; width: 300px;">
    <button onclick="sendToAider()" style="margin-left: 10px;">Send to Aider</button>
    <button onclick="listFiles()" style="margin-left: 10px;">List Files</button>
    <input id="filename_input" type="text" placeholder="Filename to run/read" style="margin-top: 10px; width: 300px;">
    <button onclick="runFile()" style="margin-left: 10px;">Run File</button>
    <button onclick="terminateProcess()" style="margin-left: 10px;">Terminate Process</button>
    <button onclick="readFile()" style="margin-left: 10px;">Read File</button>
    <button onclick="downloadFiles()" style="margin-left: 10px;">Download Files</button>
</div>
<script>
    let ws;

    function connectToAider() {
        const repoUrl = document.getElementById('repo_url').value;

        // Close existing connection if any
        if (ws) {
            ws.close();
        }

        // Create a new WebSocket connection
        ws = new WebSocket('ws://localhost:8000/ws');

        ws.onopen = function() {
            // Send initial connection parameters
            ws.send(JSON.stringify({
                repo_url: repoUrl || null
            }));

            const messagesDiv = document.getElementById('messages');
            messagesDiv.innerHTML += '<div style="color: green;">Connected to Aider' + (repoUrl ? ' with repository: ' + repoUrl : '') + '</div><br>';

            // Enable the input field
            document.getElementById('aider_input').disabled = false;
        };

        ws.onmessage = function(event) {
            // Check if the message is binary (zip file)
            if (event.data instanceof Blob) {
                // Create a download link for the zip file
                const url = window.URL.createObjectURL(event.data);
                const a = document.createElement('a');
                a.href = url;
                a.download = 'aider_files.zip';
                document.body.appendChild(a);
                a.click();
                window.URL.revokeObjectURL(url);
                document.body.removeChild(a);
                return;
            }

            const data = JSON.parse(event.data);
            const messagesDiv = document.getElementById('messages');
            if (data.type === 'output') {
                messagesDiv.innerHTML += data.content + '<br>';
            } else if (data.type === 'files') {
                messagesDiv.innerHTML += 'Files: ' + data.content.join(', ') + '<br>';
            } else if (data.type === 'file_content') {
                messagesDiv.innerHTML += 'File content:<br>' + data.content + '<br>';
            } else if (data.type === 'error') {
                messagesDiv.innerHTML += '<div style="color: red;">Error: ' + data.content + '</div><br>';
            } else if (data.type === 'status') {
                messagesDiv.innerHTML += '<div style="color: blue;">Status: ' + data.content + '</div><br>';
            }
            messagesDiv.scrollTop = messagesDiv.scrollHeight;
        };

        ws.onerror = function(error) {
            const messagesDiv = document.getElementById('messages');
            messagesDiv.innerHTML += '<div style="color: red;">WebSocket Error: ' + error.message + '</div><br>';
            document.getElementById('aider_input').disabled = true;
        };

        ws.onclose = function(event) {
            const messagesDiv = document.getElementById('messages');
            if (event.wasClean) {
                messagesDiv.innerHTML += '<div style="color: blue;">Connection closed cleanly, code=' + event.code + ' reason=' + event.reason + '</div><br>';
            } else {
                messagesDiv.innerHTML += '<div style="color: red;">Connection died unexpectedly</div><br>';
            }
            document.getElementById('aider_input').disabled = true;
        };
    }

    function sendToAider() {
        if (!ws || ws.readyState !== WebSocket.OPEN) {
            alert('Please connect to Aider first');
            return;
        }

        const input = document.getElementById('aider_input').value;
        if (input) {
            ws.send(JSON.stringify({command: 'send_to_aider', message: input}));
            document.getElementById('aider_input').value = '';
        }
    }

    function listFiles() {
        if (!ws || ws.readyState !== WebSocket.OPEN) {
            alert('Please connect to Aider first');
            return;
        }

        ws.send(JSON.stringify({command: 'list_files'}));
    }

    function readFile() {
        if (!ws || ws.readyState !== WebSocket.OPEN) {
            alert('Please connect to Aider first');
            return;
        }

        const filename = document.getElementById('filename_input').value;
        if (filename) {
            ws.send(JSON.stringify({command: 'read_file', filename: filename}));
        }
    }

    function runFile() {
        if (!ws || ws.readyState !== WebSocket.OPEN) {
            alert('Please connect to Aider first');
            return;
        }

        const filename = document.getElementById('filename_input').value;
        if (filename) {
            ws.send(JSON.stringify({command: 'run_file', filename: filename}));
        }
    }

    function terminateProcess() {
        if (!ws || ws.readyState !== WebSocket.OPEN) {
            alert('Please connect to Aider first');
            return;
        }

        ws.send(JSON.stringify({command: 'terminate_process'}));
    }

    function downloadFiles() {
        if (!ws || ws.readyState !== WebSocket.OPEN) {
            alert('Please connect to Aider first');
            return;
        }

        ws.send(JSON.stringify({command: 'download'}));
    }

    // Show connection button by default
    window.onload = function() {
        document.getElementById('aider_input').disabled = true;
    };
</script>
"""

components.html(html_code, height=500)