import { randomUUID } from 'crypto' import { writeFileSync } from 'fs' import { BrowserWindow, dialog, ipcMain } from 'electron' import { IPC } from '../../../shared/ipc-channels' import type { AiConfig, AiContextBinding, NamingConflict } from '../../../shared/types' import { BookRegistryService } from '../../services/book-registry' import { GlobalSettingsService } from '../../services/global-settings' import { AiClientService } from '../../services/ai-client.service' import { aiContextBuilder } from '../../services/ai-context-builder.service' import { buildNamingPrompt, buildNamingCorpus, collectChaptersPlain, findLocalNamingConflicts, mergeNamingConflicts, parseNamingJson } from '../../services/naming-check.service' import type { ChatMessage } from '../../services/ai-client.service' import { wrap } from '../result' const activeChats = new Map() function formatSessionMarkdown(title: string, messages: { role: string; content: string; createdAt: string }[]): string { const lines = [`# ${title || 'AI 会话'}`, ''] for (const msg of messages) { const label = msg.role === 'user' ? '用户' : msg.role === 'assistant' ? '助手' : msg.role lines.push(`## ${label}`, '', msg.content, '', `_${msg.createdAt}_`, '') } return lines.join('\n') } export function registerAiHandlers( registry: BookRegistryService, settings: GlobalSettingsService, aiClient: AiClientService ): void { ipcMain.handle( IPC.AI_SESSION_LIST, (_e, { bookId, includeArchived }: { bookId: string; includeArchived?: boolean }) => wrap(() => registry.getAiSessionRepo(bookId).list(includeArchived)) ) ipcMain.handle( IPC.AI_SESSION_CREATE, (_e, { bookId, title, model }: { bookId: string; title?: string; model?: string }) => wrap(() => { const config = settings.get().aiConfig return registry.getAiSessionRepo(bookId).create(title, model ?? config.model) }) ) ipcMain.handle( IPC.AI_SESSION_UPDATE, ( _e, { bookId, sessionId, patch }: { bookId: string sessionId: string patch: Partial<{ title: string; model: string; archived: boolean; context: AiContextBinding }> } ) => wrap(() => registry.getAiSessionRepo(bookId).update(sessionId, patch)) ) ipcMain.handle( IPC.AI_SESSION_DELETE, (_e, { bookId, sessionId }: { bookId: string; sessionId: string }) => wrap(() => registry.getAiSessionRepo(bookId).delete(sessionId)) ) ipcMain.handle( IPC.AI_MESSAGE_LIST, (_e, { bookId, sessionId }: { bookId: string; sessionId: string }) => wrap(() => registry.getAiSessionRepo(bookId).listMessages(sessionId)) ) ipcMain.handle( IPC.AI_SESSION_EXPORT, async ( _e, { bookId, sessionId, format = 'markdown' }: { bookId: string; sessionId: string; format?: 'markdown' | 'txt' } ) => wrap(async () => { const repo = registry.getAiSessionRepo(bookId) const session = repo.get(sessionId) if (!session) throw new Error('AI session not found') const messages = repo.listMessages(sessionId) const text = format === 'markdown' ? formatSessionMarkdown(session.title, messages) : messages .map((m) => `[${m.role}] ${m.createdAt}\n${m.content}`) .join('\n\n---\n\n') const ext = format === 'markdown' ? 'md' : 'txt' const { canceled, filePath } = await dialog.showSaveDialog({ defaultPath: `${session.title || 'ai-session'}.${ext}`, filters: [{ name: format === 'markdown' ? 'Markdown' : 'Text', extensions: [ext] }] }) if (canceled || !filePath) return null writeFileSync(filePath, text, 'utf8') return filePath }) ) ipcMain.handle( IPC.AI_BUILD_CONTEXT, (_e, { bookId, binding }: { bookId: string; binding: AiContextBinding }) => wrap(() => aiContextBuilder.build(binding, registry, bookId, settings.get().penName)) ) ipcMain.handle(IPC.AI_TEST_CONNECTION, (_e, { config }: { config?: AiConfig }) => wrap(async () => { const effective = config ?? settings.get().aiConfig const probe = new AiClientService(() => effective) const content = await probe.chat([{ role: 'user', content: '回复一个字:好' }]) if (!content.trim()) throw new Error('AI returned empty response') return { ok: true as const } }) ) ipcMain.handle(IPC.AI_NAMING_CHECK, (_e, { bookId }: { bookId: string }) => wrap(async () => { const settingsList = registry.getSettingRepo(bookId).list() const chapters = registry.getChapterRepo(bookId).list() const ignoreWords = settings.get().namingIgnoreWords const settingNames = settingsList.map((s) => s.name) const corpus = buildNamingCorpus(settingsList, chapters) const prompt = buildNamingPrompt(settingNames, corpus, ignoreWords) const local = findLocalNamingConflicts(settingNames, corpus, ignoreWords) if (local.length > 0) { try { const response = await Promise.race([ aiClient.chat([{ role: 'user', content: prompt }]), new Promise((_, reject) => setTimeout(() => reject(new Error('ai-timeout')), 8_000) ) ]) return mergeNamingConflicts(local, parseNamingJson(response)) } catch { return local } } const fetchAi = async (): Promise => { const response = await aiClient.chat([{ role: 'user', content: prompt }]) return parseNamingJson(response) } try { const ai = await fetchAi() return mergeNamingConflicts(local, ai) } catch { try { const response = await aiClient.chat([ { role: 'user', content: `${prompt}\n\n仅返回 JSON 数组,不要 markdown 或其它说明。` } ]) return mergeNamingConflicts(local, parseNamingJson(response)) } catch { if (local.length > 0) return local throw new Error('命名检查失败') } } }) ) ipcMain.handle(IPC.AI_ABORT, (_e, { messageId }: { messageId: string }) => wrap(() => { activeChats.get(messageId)?.abort() activeChats.delete(messageId) }) ) ipcMain.handle( IPC.AI_CHAT, ( event, { bookId, sessionId, content, prompt }: { bookId: string; sessionId: string; content: string; prompt?: string; slashCommand?: string } ) => wrap(async () => { const repo = registry.getAiSessionRepo(bookId) const session = repo.get(sessionId) if (!session) throw new Error('AI session not found') const actualPrompt = prompt ?? content repo.addMessage(sessionId, 'user', content) const history = repo.listMessages(sessionId) const { systemPrompt } = aiContextBuilder.build( session.context, registry, bookId, settings.get().penName ) const apiMessages: ChatMessage[] = [ { role: 'system', content: systemPrompt }, ...history.map((m, index) => ({ role: m.role, content: index === history.length - 1 && m.role === 'user' ? actualPrompt : m.content })) ] const assistantId = randomUUID() const controller = new AbortController() activeChats.set(assistantId, controller) const win = BrowserWindow.fromWebContents(event.sender) let partial = '' try { const full = await aiClient.chat( apiMessages, (delta) => { partial += delta win?.webContents.send(IPC.AI_STREAM_CHUNK, { messageId: assistantId, delta, done: false }) }, controller.signal ) win?.webContents.send(IPC.AI_STREAM_CHUNK, { messageId: assistantId, delta: '', done: true }) repo.addMessage(sessionId, 'assistant', full, assistantId) } catch (err) { if (controller.signal.aborted && partial) { win?.webContents.send(IPC.AI_STREAM_CHUNK, { messageId: assistantId, delta: '', done: true }) repo.addMessage(sessionId, 'assistant', partial, assistantId) return { messageId: assistantId } } throw err } finally { activeChats.delete(assistantId) } if (!session.title.trim()) { repo.update(sessionId, { title: content.slice(0, 40) }) } return { messageId: assistantId } }) ) }