bilin/src/main/ipc/handlers/ai.handler.ts

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<string, AbortController>()

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<never>((_, reject) =>
              setTimeout(() => reject(new Error('ai-timeout')), 8_000)
            )
          ])
          return mergeNamingConflicts(local, parseNamingJson(response))
        } catch {
          return local
        }
      }

      const fetchAi = async (): Promise<NamingConflict[]> => {
        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 }
      })
  )
}