bilin/src/main/services/ai-context-builder.service.ts

import { stripHtml } from './word-count'
import type {
  AiContextBinding,
  AiContextBuildResult,
  ContextPreviewItem,
  KnowledgeType
} from '../../shared/types'
import type { BookRegistryService } from './book-registry'
import { KnowledgeRepository } from '../db/repositories/knowledge.repo'

const CHAPTER_CHAR_LIMIT = 2000
const OTHER_CHAR_LIMIT = 500
const KNOWLEDGE_CHAR_LIMIT = 500
const KNOWLEDGE_TOTAL_LIMIT = 4000
const TOTAL_CHAR_LIMIT = 16 * 1024

interface BlockCandidate {
  sortKey: number
  text: string
  preview: ContextPreviewItem
  isKnowledge?: boolean
}

function truncate(text: string, max: number): string {
  if (text.length <= max) return text
  return `${text.slice(0, max)}…`
}

function knowledgeTypeLabel(type: KnowledgeType): string {
  const map: Record<KnowledgeType, string> = {
    character: '角色',
    foreshadow: '伏笔',
    location: '地点',
    relationship: '关系',
    fact: '事实'
  }
  return map[type]
}

function formatBlock(kind: ContextPreviewItem['kind'], title: string, body: string): string {
  const label =
    kind === 'chapter'
      ? '章节'
      : kind === 'outline'
        ? '大纲'
        : kind === 'setting'
          ? '设定'
          : '灵感'
  return `【${label}】${title}\n${body}`
}

function formatKnowledgeBlock(type: KnowledgeType, title: string, body: string): string {
  return `【知识·${knowledgeTypeLabel(type)}】${title}\n${body}`
}

export class AiContextBuilderService {
  build(
    binding: AiContextBinding,
    registry: BookRegistryService,
    bookId: string,
    penName: string
  ): AiContextBuildResult {
    const meta = registry.getMeta(bookId)
    if (!meta) throw new Error('Book not found')

    const candidates: BlockCandidate[] = []
    const knowledgeIds = binding.knowledgeIds ?? []

    for (const id of binding.chapterIds) {
      const chapter = registry.getChapterRepo(bookId).get(id)
      if (!chapter) continue
      const plain = stripHtml(chapter.content)
      const body = truncate(plain, CHAPTER_CHAR_LIMIT)
      candidates.push({
        sortKey: chapter.sortOrder,
        text: formatBlock('chapter', chapter.title, body),
        preview: {
          kind: 'chapter',
          id,
          title: chapter.title,
          excerpt: body,
          charCount: body.length
        }
      })
    }

    for (const id of binding.outlineIds) {
      const item = registry.getOutlineRepo(bookId).get(id)
      if (!item) continue
      const body = truncate(stripHtml(item.description), OTHER_CHAR_LIMIT)
      candidates.push({
        sortKey: item.sortOrder,
        text: formatBlock('outline', item.title, body),
        preview: {
          kind: 'outline',
          id,
          title: item.title,
          excerpt: body,
          charCount: body.length
        }
      })
    }

    for (const id of binding.settingIds) {
      const item = registry.getSettingRepo(bookId).get(id)
      if (!item) continue
      const body = truncate(stripHtml(item.description), OTHER_CHAR_LIMIT)
      candidates.push({
        sortKey: 0,
        text: formatBlock('setting', item.name, body),
        preview: {
          kind: 'setting',
          id,
          title: item.name,
          excerpt: body,
          charCount: body.length
        }
      })
    }

    for (const id of binding.inspirationIds) {
      const item = registry.getInspirationRepo(bookId).get(id)
      if (!item) continue
      const body = truncate(stripHtml(item.content), OTHER_CHAR_LIMIT)
      candidates.push({
        sortKey: Date.parse(item.updatedAt) || 0,
        text: formatBlock('inspiration', item.title, body),
        preview: {
          kind: 'inspiration',
          id,
          title: item.title,
          excerpt: body,
          charCount: body.length
        }
      })
    }

    const knowledgeRepo = new KnowledgeRepository(registry.getDb(bookId))
    const knowledgeCandidates: BlockCandidate[] = []
    for (const id of knowledgeIds) {
      const entry = knowledgeRepo.get(id)
      if (!entry || entry.status !== 'approved') continue
      const body = truncate(stripHtml(entry.content || entry.title), KNOWLEDGE_CHAR_LIMIT)
      knowledgeCandidates.push({
        sortKey: entry.importance * 1000,
        isKnowledge: true,
        text: formatKnowledgeBlock(entry.type, entry.title, body),
        preview: {
          kind: 'knowledge',
          id,
          title: entry.title,
          excerpt: body,
          charCount: body.length,
          knowledgeType: entry.type
        }
      })
    }
    knowledgeCandidates.sort((a, b) => b.sortKey - a.sortKey)

    let knowledgeChars = 0
    const trimmedKnowledge: BlockCandidate[] = []
    for (const k of knowledgeCandidates) {
      const next = knowledgeChars + k.text.length + (trimmedKnowledge.length > 0 ? 2 : 0)
      if (next > KNOWLEDGE_TOTAL_LIMIT) break
      trimmedKnowledge.push(k)
      knowledgeChars = next
    }

    candidates.push(...trimmedKnowledge)
    candidates.sort((a, b) => b.sortKey - a.sortKey)

    const selected: BlockCandidate[] = []
    let totalChars = 0
    for (const candidate of candidates) {
      const nextTotal = totalChars + candidate.text.length + (selected.length > 0 ? 2 : 0)
      if (nextTotal > TOTAL_CHAR_LIMIT) break
      selected.push(candidate)
      totalChars = nextTotal
    }

    const contextBlocks = selected.map((x) => x.text).join('\n\n')
    const systemPrompt = `你是笔临 AI 写作助手。作者笔名：${penName}。书籍：${meta.name}。
以下是与本书相关的上下文（由作者勾选）：
${contextBlocks || '（无）'}
请用简体中文回答，尊重作者的创作主权。`

    return {
      systemPrompt,
      preview: selected.map((x) => x.preview)
    }
  }
}

export const aiContextBuilder = new AiContextBuilderService()