// Initialize Socket.IO connection
const socket = io();

// DOM Elements
const problemInput = document.getElementById('problemInput');
const providerSelect = document.getElementById('providerSelect');
const maxRunsInput = document.getElementById('maxRunsInput');
const numAgentsInput = document.getElementById('numAgentsInput');
const otherPromptsInput = document.getElementById('otherPromptsInput');
const solveButton = document.getElementById('solveButton');

const statusSection = document.getElementById('statusSection');
const statusText = document.getElementById('statusText');
const sessionId = document.getElementById('sessionId');
const providerInfo = document.getElementById('providerInfo');
const agentsInfo = document.getElementById('agentsInfo');

const logSection = document.getElementById('logSection');
const logOutput = document.getElementById('logOutput');

const solutionSection = document.getElementById('solutionSection');
const solutionOutput = document.getElementById('solutionOutput');
const solutionLogPath = document.getElementById('solutionLogPath');

const errorSection = document.getElementById('errorSection');
const errorOutput = document.getElementById('errorOutput');

let currentSessionId = null;

// Event Listeners
solveButton.addEventListener('click', solveProblem);

// Socket event handlers
socket.on('connect', () => {
    console.log('Connected to server');
});

socket.on('status', (data) => {
    if (data.session_id === currentSessionId) {
        updateStatus(data.message);
    }
});

socket.on('log', (data) => {
    if (data.session_id === currentSessionId) {
        appendLog(data.line);
    }
});

socket.on('complete', (data) => {
    if (data.session_id === currentSessionId) {
        handleCompletion(data);
    }
});

socket.on('error', (data) => {
    if (data.session_id === currentSessionId) {
        handleError(data.error);
    }
});

// Functions
async function solveProblem() {
    const problem = problemInput.value.trim();

    if (!problem) {
        alert('Please enter a problem statement');
        return;
    }

    // Disable button and show status
    solveButton.disabled = true;
    solveButton.innerHTML = '<div class="spinner"></div> Solving...';

    // Hide previous results
    solutionSection.style.display = 'none';
    errorSection.style.display = 'none';

    // Show status and log sections
    statusSection.style.display = 'block';
    logSection.style.display = 'block';

    // Clear log
    logOutput.innerHTML = '';

    // Generate session ID
    currentSessionId = Date.now().toString();

    // Update info
    sessionId.textContent = currentSessionId;
    providerInfo.textContent = providerSelect.options[providerSelect.selectedIndex].text;
    agentsInfo.textContent = numAgentsInput.value;

    try {
        const response = await fetch('/api/solve', {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json',
            },
            body: JSON.stringify({
                problem: problem,
                provider: providerSelect.value,
                max_runs: parseInt(maxRunsInput.value),
                num_agents: parseInt(numAgentsInput.value),
                other_prompts: otherPromptsInput.value,
                session_id: currentSessionId
            })
        });

        const data = await response.json();

        if (!response.ok) {
            throw new Error(data.error || 'Failed to start solving');
        }

        updateStatus('Problem submitted successfully');
        appendLog(`Session ${data.session_id} started`);

    } catch (error) {
        handleError(error.message);
        solveButton.disabled = false;
        solveButton.innerHTML = `
            <svg class="button-icon" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
                <polygon points="5 3 19 12 5 21 5 3"/>
            </svg>
            Solve Problem
        `;
    }
}

function updateStatus(message) {
    statusText.textContent = message;
}

function appendLog(line) {
    if (!line.trim()) return;

    const logLine = document.createElement('div');
    logLine.className = 'log-line';

    // Add special styling for important lines
    if (line.includes('SOLUTION FOUND') || line.includes('Found a correct solution')) {
        logLine.classList.add('success');
    } else if (line.includes('ERROR') || line.includes('FAILED')) {
        logLine.classList.add('error');
    } else if (line.includes('WARNING')) {
        logLine.classList.add('warning');
    }

    logLine.textContent = line;
    logOutput.appendChild(logLine);

    // Auto-scroll to bottom
    logOutput.scrollTop = logOutput.scrollHeight;
}

function handleCompletion(data) {
    // Re-enable button
    solveButton.disabled = false;
    solveButton.innerHTML = `
        <svg class="button-icon" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
            <polygon points="5 3 19 12 5 21 5 3"/>
        </svg>
        Solve Problem
    `;

    // Update status
    if (data.solution_found) {
        updateStatus('Solution found!');

        // Show solution section
        solutionSection.style.display = 'block';
        solutionOutput.textContent = data.solution;
        solutionLogPath.textContent = `Log file: ${data.log_file}`;

        // Add success message to log
        appendLog('='.repeat(60));
        appendLog('SOLUTION FOUND!');
        appendLog('='.repeat(60));
    } else {
        updateStatus('Completed - No solution found');

        // Add completion message to log
        appendLog('='.repeat(60));
        appendLog('Process completed without finding a solution');
        appendLog('Check the log above for partial results or insights');
        appendLog('='.repeat(60));
    }
}

function handleError(errorMessage) {
    // Re-enable button
    solveButton.disabled = false;
    solveButton.innerHTML = `
        <svg class="button-icon" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
            <polygon points="5 3 19 12 5 21 5 3"/>
        </svg>
        Solve Problem
    `;

    // Update status
    updateStatus('Error occurred');

    // Show error section
    errorSection.style.display = 'block';
    errorOutput.textContent = errorMessage;

    // Add to log
    appendLog('ERROR: ' + errorMessage);
}

// Add keyboard shortcut (Ctrl+Enter or Cmd+Enter to solve)
problemInput.addEventListener('keydown', (e) => {
    if ((e.ctrlKey || e.metaKey) && e.key === 'Enter') {
        e.preventDefault();
        solveProblem();
    }
});

// Add sample problems
const sampleProblems = [
    `Let n be a positive integer. There are n red cards and n blue cards. Initially, every red card has the real number 0 written on it, and every blue card has the real number 1 written on it.
An operation consists of selecting one red card and one blue card such that the number x on the red card is strictly less than the number y on the blue card, erasing both numbers, and replacing them both with the average (x+y)/2.
Find the minimum positive integer n such that it is possible, using a finite sequence of these operations, to make the sum of the numbers on the n red cards strictly greater than 100.`,

    `Find all prime numbers p for which there exist non-negative integers x, y, and z such that p = x^3 + y^3 + z^3.`,

    `Determine all positive integers n for which there exists a permutation (a₁, a₂, ..., aₙ) of (1, 2, ..., n) such that for each i ∈ {1, 2, ..., n}, the sum a₁ + a₂ + ... + aᵢ is divisible by i.`
];

// Add sample problem selector (optional feature)
const createSampleSelector = () => {
    const sampleSelector = document.createElement('div');
    sampleSelector.className = 'control-group full-width';
    sampleSelector.innerHTML = `
        <label class="control-label">
            <svg class="label-icon" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
                <path d="M12 2L2 7l10 5 10-5-10-5zM2 17l10 5 10-5M2 12l10 5 10-5"/>
            </svg>
            Or try a sample problem:
        </label>
        <select id="sampleSelect" class="control-select">
            <option value="">-- Select a sample problem --</option>
            <option value="0">IMO-style: Red and Blue Cards</option>
            <option value="1">Number Theory: Prime Sums</option>
            <option value="2">Combinatorics: Permutation Divisibility</option>
        </select>
    `;

    const controls = document.querySelector('.controls');
    controls.insertBefore(sampleSelector, controls.firstChild);

    document.getElementById('sampleSelect').addEventListener('change', (e) => {
        const index = parseInt(e.target.value);
        if (!isNaN(index) && index >= 0 && index < sampleProblems.length) {
            problemInput.value = sampleProblems[index];
        }
    });
};

// Initialize sample selector when page loads
window.addEventListener('load', createSampleSelector);