wqb-server/static/brain.js

/**
 * BRAIN API 集成模块
 * 处理 WorldQuant BRAIN 的身份验证、运算符和数据字段
 * 现在使用本地 Python 代理服务器来绕过 CORS 限制
 */

// BRAIN 会话和数据存储
let brainSession = null;
let brainOperators = null;
let brainDataFields = null;
let brainSessionId = localStorage.getItem('brain_session_id');

// Flask 应用 API 端点
const PROXY_BASE = '';

// 打开 BRAIN 登录模态框
function openBrainLoginModal() {
    const modal = document.getElementById('brainLoginModal');
    const statusDiv = document.getElementById('brainLoginStatus');
    statusDiv.innerHTML = '';
    statusDiv.className = 'login-status';

    // 清除之前的输入
    document.getElementById('brainUsername').value = '';
    document.getElementById('brainPassword').value = '';

    modal.style.display = 'block';
    document.getElementById('brainUsername').focus();
}

// 关闭 BRAIN 登录模态框
function closeBrainLoginModal() {
    const modal = document.getElementById('brainLoginModal');
    const loginBtn = document.getElementById('loginBtn');

    // 如果登录正在进行中，不允许关闭
    if (loginBtn.disabled) {
        return;
    }

    modal.style.display = 'none';
}

// 通过代理服务器与 BRAIN 进行身份验证
async function authenticateBrain() {
    const username = document.getElementById('brainUsername').value.trim();
    const password = document.getElementById('brainPassword').value;
    const statusDiv = document.getElementById('brainLoginStatus');
    const loginBtn = document.getElementById('loginBtn');
    const spinner = document.getElementById('loginSpinner');
    const modal = document.getElementById('brainLoginModal');

    if (!username || !password) {
        showLoginStatus('请输入用户名和密码。', 'error');
        return;
    }

    // 禁用所有输入框和按钮
    document.getElementById('brainUsername').disabled = true;
    document.getElementById('brainPassword').disabled = true;
    document.getElementById('cancelBtn').disabled = true;
    loginBtn.disabled = true;
    loginBtn.textContent = '连接中...';

    // 显示加载指示器
    spinner.style.display = 'block';

    // 禁用模态框关闭
    modal.querySelector('.close').style.display = 'none';

    // 显示加载状态
    showLoginStatus('正在连接到代理服务器...', 'loading');

    try {
        showLoginStatus('正在与 BRAIN 进行身份验证...', 'loading');

        // 通过代理服务器进行身份验证
        const authResponse = await fetch(`${PROXY_BASE}/api/authenticate`, {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json'
            },
            body: JSON.stringify({
                username: username,
                password: password
            })
        });

        if (!authResponse.ok) {
            const errorData = await authResponse.json();
            throw new Error(errorData.error || '身份验证失败');
        }

        const authData = await authResponse.json();
        brainSessionId = authData.session_id;
        brainSession = { authenticated: true, username: username };

        // 将会话 ID 存储在 localStorage 中供其他页面使用
        localStorage.setItem('brain_session_id', brainSessionId);

        // 立即获取运算符以支持 "Op" 按钮功能
        showLoginStatus('正在加载运算符...', 'loading');
        brainOperators = await getUserOperators();

        // 更新 UI 显示已连接状态
        updateConnectedState();
        showLoginStatus(`连接成功！已加载 ${brainOperators.length} 个运算符。`, 'success');

        // 禁用按钮以防止进一步点击
        loginBtn.disabled = true;
        document.getElementById('brainUsername').disabled = true;
        document.getElementById('brainPassword').disabled = true;

        // 短暂延迟后关闭模态框
        setTimeout(() => {
            // 在关闭前重新启用所有控件
            document.getElementById('brainUsername').disabled = false;
            document.getElementById('brainPassword').disabled = false;
            document.getElementById('cancelBtn').disabled = false;
            loginBtn.disabled = false;
            loginBtn.textContent = '连接';
            spinner.style.display = 'none';
            modal.querySelector('.close').style.display = 'block';

            closeBrainLoginModal();
        }, 1500);

    } catch (error) {
        console.error('BRAIN 身份验证失败:', error);
        showLoginStatus(`连接失败: ${error.message}`, 'error');
        brainSession = null;
        brainSessionId = null;
    } finally {
        // 重新启用所有控件
        document.getElementById('brainUsername').disabled = false;
        document.getElementById('brainPassword').disabled = false;
        document.getElementById('cancelBtn').disabled = false;
        loginBtn.disabled = false;
        loginBtn.textContent = '连接';
        spinner.style.display = 'none';
        modal.querySelector('.close').style.display = 'block';
    }
}

// 通过代理服务器获取用户运算符
async function getUserOperators() {
    if (!brainSession || !brainSessionId) {
        throw new Error('未通过 BRAIN 身份验证');
    }

    try {
        const response = await fetch(`${PROXY_BASE}/api/operators`, {
            method: 'GET',
            headers: {
                'Session-ID': brainSessionId
            }
        });

        if (!response.ok) {
            const errorData = await response.json();
            throw new Error(errorData.error || '获取运算符失败');
        }

        const operators = await response.json();
        console.log(`从 BRAIN API 接收到 ${operators.length} 个运算符`);

        // 记录类别以验证我们拥有所有运算符类型
        const categories = [...new Set(operators.map(op => op.category))].sort();
        console.log(`运算符类别: ${categories.join(', ')}`);

        return operators;

    } catch (error) {
        console.error('获取运算符失败:', error);
        throw error;
    }
}

// 通过代理服务器获取数据字段
async function getDataFields(region = 'USA', delay = 1, universe = 'TOP3000', datasetId = 'fundamental6') {
    if (!brainSession || !brainSessionId) {
        throw new Error('未通过 BRAIN 身份验证');
    }

    try {
        const params = new URLSearchParams({
            region: region,
            delay: delay.toString(),
            universe: universe,
            dataset_id: datasetId
        });

        const response = await fetch(`${PROXY_BASE}/api/datafields?${params}`, {
            method: 'GET',
            headers: {
                'Session-ID': brainSessionId
            }
        });

        if (!response.ok) {
            const errorData = await response.json();
            throw new Error(errorData.error || '获取数据字段失败');
        }

        return await response.json();

    } catch (error) {
        console.error('获取数据字段失败:', error);
        throw error;
    }
}

// 更新 UI 显示已连接状态
function updateConnectedState() {
    const connectBtn = document.getElementById('connectToBrain');
    connectBtn.textContent = '已连接到 BRAIN';
    connectBtn.className = 'btn btn-brain connected';

    // 在语法错误区域显示连接信息
    const errorsDiv = document.getElementById('grammarErrors');
    errorsDiv.innerHTML = `<div class="success-message">
        ✓ 已成功连接到 WorldQuant BRAIN<br>
        <strong>用户名:</strong> ${brainSession.username}<br>
        <strong>已加载运算符:</strong> ${brainOperators ? brainOperators.length : 0}<br>
        <em>数据字段将在需要时加载。</em>
    </div>`;

    // 5 秒后自动隐藏消息
    setTimeout(() => {
        if (errorsDiv.innerHTML.includes('已成功连接')) {
            errorsDiv.innerHTML = '';
        }
    }, 5000);
}

// 显示登录状态消息
function showLoginStatus(message, type) {
    const statusDiv = document.getElementById('brainLoginStatus');
    statusDiv.textContent = message;
    statusDiv.className = `login-status ${type}`;
}

// 检查是否已连接到 BRAIN
function isConnectedToBrain() {
    return brainSession !== null && brainSessionId !== null;
}

// 获取所有可用运算符（按需获取）
async function getAllOperators() {
    if (!brainOperators && isConnectedToBrain()) {
        try {
            brainOperators = await getUserOperators();
        } catch (error) {
            console.error('按需获取运算符失败:', error);
            return [];
        }
    }
    return brainOperators || [];
}

// 同步获取已加载的运算符（用于 UI 组件）
function getLoadedOperators() {
    return brainOperators || [];
}

// 获取所有可用数据字段（按需获取）
async function getAllDataFields() {
    if (!brainDataFields && isConnectedToBrain()) {
        try {
            brainDataFields = await getDataFields();
        } catch (error) {
            console.error('按需获取数据字段失败:', error);
            return [];
        }
    }
    return brainDataFields || [];
}

// 按类别获取运算符（支持按需加载）
async function getOperatorsByCategory(category) {
    const operators = await getAllOperators();
    return operators.filter(op => op.category === category);
}

// 搜索运算符（支持按需加载）
async function searchOperators(searchTerm) {
    const operators = await getAllOperators();
    const term = searchTerm.toLowerCase();
    return operators.filter(op =>
        op.name.toLowerCase().includes(term) ||
        op.category.toLowerCase().includes(term)
    );
}

// 搜索数据字段（支持按需加载）
async function searchDataFields(searchTerm) {
    const dataFields = await getAllDataFields();
    const term = searchTerm.toLowerCase();
    return dataFields.filter(field =>
        field.id.toLowerCase().includes(term) ||
        field.description.toLowerCase().includes(term)
    );
}

// 从 BRAIN 注销
async function logoutFromBrain() {
    if (brainSessionId) {
        try {
            await fetch(`${PROXY_BASE}/api/logout`, {
                method: 'POST',
                headers: {
                    'Session-ID': brainSessionId
                }
            });
        } catch (error) {
            console.warn('从代理服务器注销失败:', error);
        }
    }

    // 清除本地会话数据
    brainSession = null;
    brainSessionId = null;
    brainOperators = null;
    brainDataFields = null;

    // 清除 localStorage
    localStorage.removeItem('brain_session_id');

    // 更新 UI
    const connectBtn = document.getElementById('connectToBrain');
    connectBtn.textContent = '连接到 BRAIN';
    connectBtn.className = 'btn btn-brain';
}

// 在页面加载时检查会话有效性
async function checkSessionValidity() {
    if (brainSessionId) {
        try {
            const response = await fetch(`${PROXY_BASE}/api/status`, {
                method: 'GET',
                headers: {
                    'Session-ID': brainSessionId
                }
            });

            if (response.ok) {
                const data = await response.json();
                if (data.valid) {
                    brainSession = { authenticated: true, username: data.username };
                    // 更新 UI 显示已连接状态
                    updateConnectedState();
                } else {
                    // 会话已过期，清除它
                    localStorage.removeItem('brain_session_id');
                    brainSessionId = null;
                }
            }
        } catch (error) {
            console.warn('检查会话有效性失败:', error);
        }
    }
}

// 在页面加载时初始化
document.addEventListener('DOMContentLoaded', checkSessionValidity);

// 导出函数供其他模块使用
window.brainAPI = {
    openBrainLoginModal,
    closeBrainLoginModal,
    authenticateBrain,
    isConnectedToBrain,
    getAllOperators,
    getAllDataFields,
    getDataFields,
    getOperatorsByCategory,
    searchOperators,
    searchDataFields,
    logoutFromBrain,
    getLoadedOperators
};