import { readFile } from 'fs/promises'
import matter from 'gray-matter'

export interface DocMeta {
  title: string
  description?: string
  tags?: string[]
  date?: string
}

export interface ParsedDoc {
  path: string
  meta: DocMeta
  content: string
  headings: string[]
}

export async function parseMarkdown(filePath: string): Promise<ParsedDoc> {
  const raw = await readFile(filePath, 'utf-8')
  const { data, content } = matter(raw)

  // 提取标题
  const headings = content
    .split('\n')
    .filter((line) => line.startsWith('#'))
    .map((line) => line.replace(/^#+\s*/, ''))

  // 提取 meta
  const meta: DocMeta = {
    title: data.title || headings[0] || filePath.split('/').pop() || 'Untitled',
    description: data.description,
    tags: data.tags,
    date: data.date?.toString(),
  }

  return { path: filePath, meta, content, headings }
}

export function extractSummary(content: string, maxLength = 200): string {
  // 移除 Markdown 语法
  const plain = content
    .replace(/^#+\s+/gm, '') // 标题
    .replace(/\*\*|__/g, '') // 粗体
    .replace(/\*|_/g, '') // 斜体
    .replace(/`{1,3}[^`]*`{1,3}/g, '') // 代码
    .replace(/\[([^\]]+)\]\([^)]+\)/g, '$1') // 链接
    .replace(/\n+/g, ' ') // 换行
    .trim()

  if (plain.length <= maxLength) return plain
  return plain.slice(0, maxLength).trim() + '...'
}

import { glob } from 'glob'
import { parseMarkdown, type ParsedDoc } from './markdown.js'

export interface SearchResult {
  path: string
  title: string
  matches: string[]
  score: number
}

export class DocIndex {
  private docs: ParsedDoc[] = []
  private baseDir: string

  constructor(baseDir: string) {
    this.baseDir = baseDir
  }

  async build(): Promise<void> {
    const files = await glob('**/*.md', { cwd: this.baseDir, absolute: true })

    this.docs = await Promise.all(files.map((file) => parseMarkdown(file)))

    console.error(`Indexed ${this.docs.length} documents`)
  }

  search(query: string, limit = 10): SearchResult[] {
    const queryLower = query.toLowerCase()
    const queryTerms = queryLower.split(/\s+/).filter((t) => t.length > 1)

    const results: SearchResult[] = []

    for (const doc of this.docs) {
      const contentLower = doc.content.toLowerCase()
      const titleLower = doc.meta.title.toLowerCase()

      let score = 0
      const matches: string[] = []

      // 标题匹配（高权重）
      for (const term of queryTerms) {
        if (titleLower.includes(term)) {
          score += 10
        }
      }

      // 内容匹配
      for (const term of queryTerms) {
        const regex = new RegExp(term, 'gi')
        const contentMatches = contentLower.match(regex)
        if (contentMatches) {
          score += contentMatches.length

          // 提取上下文
          const lines = doc.content.split('\n')
          for (const line of lines) {
            if (line.toLowerCase().includes(term) && matches.length < 3) {
              matches.push(line.trim().slice(0, 100))
            }
          }
        }
      }

      if (score > 0) {
        results.push({
          path: doc.path.replace(this.baseDir, '').replace(/^\//, ''),
          title: doc.meta.title,
          matches: [...new Set(matches)], // 去重
          score,
        })
      }
    }

    return results.sort((a, b) => b.score - a.score).slice(0, limit)
  }

  getDoc(relativePath: string): ParsedDoc | undefined {
    const fullPath = `${this.baseDir}/${relativePath}`
    return this.docs.find((d) => d.path === fullPath)
  }

  getAllDocs(): ParsedDoc[] {
    return this.docs
  }

  getStats() {
    return {
      totalDocs: this.docs.length,
      totalWords: this.docs.reduce(
        (sum, doc) => sum + doc.content.split(/\s+/).length,
        0
      ),
      tags: [...new Set(this.docs.flatMap((d) => d.meta.tags || []))],
    }
  }
}

import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { z } from 'zod'
import { DocIndex } from './utils/search.js'

export function registerTools(server: McpServer, index: DocIndex) {
  // 搜索文档
  server.registerTool(
    'search_docs',
    {
      title: '搜索文档',
      description: '在文档库中搜索包含指定关键词的内容',
      inputSchema: {
        query: z.string().min(1).describe('搜索关键词'),
        limit: z
          .number()
          .int()
          .min(1)
          .max(20)
          .default(5)
          .describe('返回结果数量'),
      },
      outputSchema: {
        results: z.array(
          z.object({
            path: z.string(),
            title: z.string(),
            matches: z.array(z.string()),
            score: z.number(),
          })
        ),
      },
    },
    async ({ query, limit }) => {
      const results = index.search(query, limit)

      if (results.length === 0) {
        return {
          content: [{ type: 'text', text: `未找到与 "${query}" 相关的文档` }],
          structuredContent: { results: [] },
        }
      }

      const text = results
        .map(
          (r, i) =>
            `${i + 1}. **${r.title}** (${r.path})\n   ${r.matches[0] || ''}`
        )
        .join('\n\n')

      return {
        content: [
          {
            type: 'text',
            text: `找到 ${results.length} 个相关文档：\n\n${text}`,
          },
        ],
        structuredContent: { results },
      }
    }
  )

  // 获取文档结构
  server.registerTool(
    'get_doc_structure',
    {
      title: '获取文档结构',
      description: '获取指定文档的标题结构（目录）',
      inputSchema: {
        path: z.string().describe('文档路径（相对路径）'),
      },
    },
    async ({ path }) => {
      const doc = index.getDoc(path)

      if (!doc) {
        return {
          content: [{ type: 'text', text: `文档不存在：${path}` }],
          isError: true,
        }
      }

      const structure = doc.headings.map((h, i) => `${i + 1}. ${h}`).join('\n')

      return {
        content: [
          {
            type: 'text',
            text: `## ${doc.meta.title}\n\n**目录结构：**\n${structure}`,
          },
        ],
      }
    }
  )
}

import {
  McpServer,
  ResourceTemplate,
} from '@modelcontextprotocol/sdk/server/mcp.js'
import { DocIndex } from './utils/search.js'
import { extractSummary } from './utils/markdown.js'

export function registerResources(server: McpServer, index: DocIndex) {
  // 文档索引
  server.registerResource(
    'docs-index',
    'docs://index',
    {
      title: '文档索引',
      description: '文档库的索引和统计信息',
      mimeType: 'application/json',
    },
    async (uri) => {
      const stats = index.getStats()
      const docs = index.getAllDocs().map((d) => ({
        path: d.path,
        title: d.meta.title,
        summary: extractSummary(d.content),
      }))

      return {
        contents: [
          {
            uri: uri.href,
            mimeType: 'application/json',
            text: JSON.stringify({ stats, docs }, null, 2),
          },
        ],
      }
    }
  )

  // 单个文档
  server.registerResource(
    'doc-content',
    new ResourceTemplate('docs://{path*}', {
      list: async () => {
        const docs = index.getAllDocs()
        return {
          resources: docs.map((d) => ({
            uri: `docs://${d.path.split('/').slice(-1)[0]}`,
            name: d.meta.title,
            description: d.meta.description || extractSummary(d.content, 100),
            mimeType: 'text/markdown',
          })),
        }
      },
    }),
    {
      title: '文档内容',
      description: '访问具体的文档内容',
    },
    async (uri, { path }) => {
      // 尝试多种路径匹配
      let doc = index.getDoc(path)
      if (!doc) {
        doc = index.getDoc(`${path}.md`)
      }
      if (!doc) {
        // 尝试文件名匹配
        const docs = index.getAllDocs()
        doc = docs.find(
          (d) => d.path.endsWith(`/${path}`) || d.path.endsWith(`/${path}.md`)
        )
      }

      if (!doc) {
        return {
          contents: [
            {
              uri: uri.href,
              mimeType: 'text/plain',
              text: `文档不存在：${path}`,
            },
          ],
        }
      }

      const header = [
        `# ${doc.meta.title}`,
        doc.meta.description ? `\n> ${doc.meta.description}` : '',
        doc.meta.tags?.length ? `\n**标签：** ${doc.meta.tags.join(', ')}` : '',
        '\n---\n',
      ]
        .filter(Boolean)
        .join('')

      return {
        contents: [
          {
            uri: uri.href,
            mimeType: 'text/markdown',
            text: header + doc.content,
          },
        ],
      }
    }
  )
}

import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { z } from 'zod'
import { DocIndex } from './utils/search.js'
import { extractSummary } from './utils/markdown.js'

export function registerPrompts(server: McpServer, index: DocIndex) {
  // 文档摘要
  server.registerPrompt(
    'summarize',
    {
      title: '文档摘要',
      description: '生成指定文档的摘要',
      argsSchema: {
        path: z.string().describe('文档路径'),
      },
    },
    ({ path }) => {
      const doc = index.getDoc(path)

      if (!doc) {
        return {
          messages: [
            {
              role: 'user',
              content: { type: 'text', text: `文档不存在：${path}` },
            },
          ],
        }
      }

      return {
        messages: [
          {
            role: 'user',
            content: {
              type: 'text',
              text: `请为以下文档生成一个简洁的摘要（不超过 200 字）：

# ${doc.meta.title}

${doc.content}

要求：
1. 概括文档的主要内容
2. 列出关键要点（3-5 个）
3. 说明目标读者`,
            },
          },
        ],
      }
    }
  )

  // 概念解释
  server.registerPrompt(
    'explain',
    {
      title: '概念解释',
      description: '解释文档中的特定概念',
      argsSchema: {
        concept: z.string().describe('要解释的概念'),
        context: z.string().optional().describe('上下文（可选）'),
      },
    },
    ({ concept, context }) => {
      // 搜索相关文档
      const results = index.search(concept, 3)

      let docContext = ''
      if (results.length > 0) {
        docContext = results
          .map((r) => {
            const doc = index.getDoc(r.path)
            return doc
              ? `## ${doc.meta.title}\n${extractSummary(doc.content, 500)}`
              : ''
          })
          .filter(Boolean)
          .join('\n\n')
      }

      return {
        messages: [
          {
            role: 'user',
            content: {
              type: 'text',
              text: `请解释「${concept}」这个概念。

${docContext ? `以下是相关文档内容供参考：\n\n${docContext}\n\n` : ''}
${context ? `额外上下文：${context}\n\n` : ''}
请提供：
1. 简明定义
2. 核心要点
3. 实际应用示例
4. 常见误区（如有）`,
            },
          },
        ],
      }
    }
  )

  // 文档问答
  server.registerPrompt(
    'qa',
    {
      title: '文档问答',
      description: '基于文档内容回答问题',
      argsSchema: {
        question: z.string().describe('你的问题'),
      },
    },
    ({ question }) => {
      // 搜索相关内容
      const results = index.search(question, 5)

      let context = ''
      if (results.length > 0) {
        context = results
          .map((r) => {
            const doc = index.getDoc(r.path)
            if (!doc) return ''
            return `### ${doc.meta.title}\n${doc.content.slice(0, 1000)}${doc.content.length > 1000 ? '...' : ''}`
          })
          .filter(Boolean)
          .join('\n\n---\n\n')
      }

      return {
        messages: [
          {
            role: 'user',
            content: {
              type: 'text',
              text: `请根据以下文档内容回答我的问题。

**问题：** ${question}

${context ? `**相关文档：**\n\n${context}` : '**注意：** 未找到相关文档，请基于你的知识回答。'}

请注意：
1. 如果文档中有答案，请引用具体内容
2. 如果文档中没有直接答案，请说明
3. 可以补充文档之外的相关知识`,
            },
          },
        ],
      }
    }
  )
}

import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { DocIndex } from './utils/search.js'
import { registerTools } from './tools.js'
import { registerResources } from './resources.js'
import { registerPrompts } from './prompts.js'

export async function createServer(docsDir: string): Promise<McpServer> {
  // 创建文档索引
  const index = new DocIndex(docsDir)
  await index.build()

  // 创建 MCP Server
  const server = new McpServer({
    name: 'docs-assistant',
    version: '1.0.0',
  })

  // 注册功能
  registerTools(server, index)
  registerResources(server, index)
  registerPrompts(server, index)

  return server
}

#!/usr/bin/env node
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'
import { createServer } from './server.js'
import { resolve } from 'path'

async function main() {
  // 从命令行参数或环境变量获取文档目录
  const docsDir = resolve(process.argv[2] || process.env.DOCS_DIR || './docs')

  console.error(`Docs directory: ${docsDir}`)

  try {
    const server = await createServer(docsDir)
    const transport = new StdioServerTransport()

    await server.connect(transport)
    console.error('MCP Docs Assistant is running')
  } catch (error) {
    console.error('Failed to start server:', error)
    process.exit(1)
  }
}

main()