/**
 * AI Handler
 *
 * Handles OpenAI-compatible API streaming for web clients.
 */
import { WebSocket } from 'ws';
import OpenAI from 'openai';
// Model configs per session
const modelConfigs = new Map();
// Get OpenAI client
function getOpenAIClient(config) {
    let baseURL = config.baseUrl;
    if (baseURL && baseURL.endsWith('/')) {
        baseURL = baseURL.slice(0, -1);
    }
    return new OpenAI({
        apiKey: config.apiKey,
        baseURL: baseURL,
        dangerouslyAllowBrowser: true,
    });
}
// AI handler
export function setupAIHandler(ws, sessionId) {
    return (msg) => {
        const { type, id, requestId, method, args } = msg;
        if (method === 'ai:chatStream') {
            handleChatStream(ws, requestId, args, sessionId);
            return;
        }
        switch (method) {
            case 'ai:setApiKey':
                handleSetApiKey(ws, id, args);
                break;
            case 'ai:hasApiKey':
                handleHasApiKey(ws, id, sessionId);
                break;
            case 'ai:setModelConfig':
                handleSetModelConfig(ws, id, args, sessionId);
                break;
            case 'ai:getModelConfig':
                handleGetModelConfig(ws, id, sessionId);
                break;
            case 'ai:complete':
                handleComplete(ws, id, args, sessionId);
                break;
            case 'ai:explain':
                handleExplain(ws, id, args, sessionId);
                break;
            case 'ai:optimize':
                handleOptimize(ws, id, args, sessionId);
                break;
            case 'ai:refactor':
                handleRefactor(ws, id, args, sessionId);
                break;
            case 'ai:generate':
                handleGenerate(ws, id, args, sessionId);
                break;
            default:
                ws.send(JSON.stringify({ type: 'error', id, error: `Unknown ai method: ${method}` }));
        }
    };
}
function handleSetApiKey(ws, id, args) {
    ws.send(JSON.stringify({ type: 'response', id, result: { success: true } }));
}
function handleHasApiKey(ws, id, sessionId) {
    const config = modelConfigs.get(sessionId);
    ws.send(JSON.stringify({ type: 'response', id, result: { hasKey: !!config?.apiKey } }));
}
function handleSetModelConfig(ws, id, args, sessionId) {
    const config = args[0];
    modelConfigs.set(sessionId, config);
    ws.send(JSON.stringify({ type: 'response', id, result: { success: true } }));
}
function handleGetModelConfig(ws, id, sessionId) {
    const config = modelConfigs.get(sessionId);
    ws.send(JSON.stringify({ type: 'response', id, result: { config } }));
}
async function handleComplete(ws, id, args, sessionId) {
    const config = modelConfigs.get(sessionId);
    if (!config || !config.apiKey) {
        ws.send(JSON.stringify({ type: 'error', id, error: 'API key not set' }));
        return;
    }
    const prefix = args[0];
    const suffix = args[1];
    const language = args[2];
    try {
        const client = getOpenAIClient(config);
        const response = await client.chat.completions.create({
            model: config.modelId,
            max_tokens: config.maxTokens || 256,
            messages: [
                { role: 'system', content: `You are a code completion assistant. Complete the code at the cursor position. Language: ${language}` },
                { role: 'user', content: `${prefix}<CURSOR>${suffix}` },
            ],
            temperature: config.temperature || 0.2,
        });
        ws.send(JSON.stringify({ type: 'response', id, result: { completion: response.choices[0]?.message?.content || '' } }));
    }
    catch (error) {
        ws.send(JSON.stringify({ type: 'error', id, error: `Completion failed: ${error.message}` }));
    }
}
async function handleExplain(ws, id, args, sessionId) {
    const config = modelConfigs.get(sessionId);
    if (!config || !config.apiKey) {
        ws.send(JSON.stringify({ type: 'error', id, error: 'API key not set' }));
        return;
    }
    const code = args[0];
    const language = args[1];
    try {
        const client = getOpenAIClient(config);
        const response = await client.chat.completions.create({
            model: config.modelId,
            max_tokens: config.maxTokens || 2048,
            messages: [{ role: 'user', content: `Explain this ${language} code:\n\n${code}` }],
            temperature: config.temperature || 0.5,
        });
        ws.send(JSON.stringify({ type: 'response', id, result: { explanation: response.choices[0]?.message?.content || '' } }));
    }
    catch (error) {
        ws.send(JSON.stringify({ type: 'error', id, error: `Explanation failed: ${error.message}` }));
    }
}
async function handleOptimize(ws, id, args, sessionId) {
    const config = modelConfigs.get(sessionId);
    if (!config || !config.apiKey) {
        ws.send(JSON.stringify({ type: 'error', id, error: 'API key not set' }));
        return;
    }
    const code = args[0];
    const language = args[1];
    try {
        const client = getOpenAIClient(config);
        const response = await client.chat.completions.create({
            model: config.modelId,
            max_tokens: config.maxTokens || 2048,
            messages: [{ role: 'user', content: `Optimize this ${language} code. Only return the optimized code:\n\n${code}` }],
            temperature: config.temperature || 0.2,
        });
        ws.send(JSON.stringify({ type: 'response', id, result: { optimizedCode: response.choices[0]?.message?.content || code } }));
    }
    catch (error) {
        ws.send(JSON.stringify({ type: 'error', id, error: `Optimization failed: ${error.message}` }));
    }
}
async function handleRefactor(ws, id, args, sessionId) {
    const config = modelConfigs.get(sessionId);
    if (!config || !config.apiKey) {
        ws.send(JSON.stringify({ type: 'error', id, error: 'API key not set' }));
        return;
    }
    const code = args[0];
    const language = args[1];
    const instructions = args[2];
    try {
        const client = getOpenAIClient(config);
        const response = await client.chat.completions.create({
            model: config.modelId,
            max_tokens: config.maxTokens || 4096,
            messages: [{ role: 'user', content: `Refactor this ${language} code to ${instructions}. Only return the refactored code:\n\n${code}` }],
            temperature: config.temperature || 0.3,
        });
        ws.send(JSON.stringify({ type: 'response', id, result: { refactoredCode: response.choices[0]?.message?.content || code } }));
    }
    catch (error) {
        ws.send(JSON.stringify({ type: 'error', id, error: `Refactoring failed: ${error.message}` }));
    }
}
async function handleGenerate(ws, id, args, sessionId) {
    const config = modelConfigs.get(sessionId);
    if (!config || !config.apiKey) {
        ws.send(JSON.stringify({ type: 'error', id, error: 'API key not set' }));
        return;
    }
    const description = args[0];
    const language = args[1];
    try {
        const client = getOpenAIClient(config);
        const response = await client.chat.completions.create({
            model: config.modelId,
            max_tokens: config.maxTokens || 4096,
            messages: [{ role: 'user', content: `Generate ${language} code for: ${description}. Only return the code.` }],
            temperature: config.temperature || 0.5,
        });
        ws.send(JSON.stringify({ type: 'response', id, result: { generatedCode: response.choices[0]?.message?.content || '' } }));
    }
    catch (error) {
        ws.send(JSON.stringify({ type: 'error', id, error: `Code generation failed: ${error.message}` }));
    }
}
async function handleChatStream(ws, requestId, args, sessionId) {
    const messages = args[0] || [];
    const options = args[1] || {};
    const config = modelConfigs.get(sessionId);
    if (!config || !config.apiKey) {
        ws.send(JSON.stringify({ type: 'streamError', requestId, error: '请先配置AI模型和API密钥' }));
        return;
    }
    try {
        const client = getOpenAIClient(config);
        // Format messages
        const formattedMessages = messages.map((msg) => ({
            role: msg.role === 'assistant' ? 'assistant' : 'user',
            content: msg.content,
        }));
        const apiMessages = options.system
            ? [{ role: 'system', content: options.system }, ...formattedMessages]
            : formattedMessages;
        ws.send(JSON.stringify({ type: 'streamStart', requestId }));
        const stream = await client.chat.completions.create({
            model: config.modelId,
            messages: apiMessages,
            max_tokens: config.maxTokens || 4096,
            temperature: config.temperature ?? 0.7,
            stream: true,
        });
        for await (const chunk of stream) {
            const content = chunk.choices[0]?.delta?.content;
            if (content && ws.readyState === WebSocket.OPEN) {
                ws.send(JSON.stringify({ type: 'streamData', requestId, data: content }));
            }
        }
        if (ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({ type: 'streamEnd', requestId }));
        }
    }
    catch (error) {
        let errorMessage = error.message || 'Unknown error';
        if (error.status === 404) {
            errorMessage = `模型 "${config.modelId}" 不存在`;
        }
        else if (error.status === 401) {
            errorMessage = `API Key 无效或已过期`;
        }
        else if (error.status === 429) {
            errorMessage = `请求过于频繁或额度不足`;
        }
        else if (error.code === 'ENOTFOUND' || error.code === 'ECONNREFUSED') {
            errorMessage = `无法连接到服务器`;
        }
        if (ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({ type: 'streamError', requestId, error: errorMessage }));
        }
    }
}
//# sourceMappingURL=ai.js.map