1import { mkdirSync, readFileSync, writeFileSync } from "fs"
2import { dirname, join } from "path"
3import MarkdownIt from "markdown-it"
4 
5type BlockKind =
6  | "heading"
7  | "thematicBreak"
8  | "html"
9  | "code"
10  | "flow"
11 
12interface Block {
13  kind: BlockKind
14  md: string
15  words: number
16}
17 
18interface Chunk {
19  blocks: Block[]
20  words: number
21}
22 
23export interface ChunkCliOptions {
24  file: string
25  maxWords: number
26  outputDir: string
27}
28 
29export interface ChunkResult {
30  source: string
31  chunks: number
32  output_dir: string
33  frontmatter: boolean
34  words_per_chunk: number[]
35}
36 
37const parser = new MarkdownIt({ html: true })
38 
39export function formatChunkUsage(command: string): string {
40  return `Usage: ${command} <file> [--max-words 5000] [--output-dir <dir>]`
41}
42 
43export function runChunkCli(args: string[], command = "chunk.ts"): number {
44  const parsed = parseChunkCliArgs(args)
45 
46  if ("help" in parsed) {
47    console.log(formatChunkUsage(command))
48    return 0
49  }
50 
51  if ("error" in parsed) {
52    console.error(parsed.error)
53    console.error(formatChunkUsage(command))
54    return 1
55  }
56 
57  const result = chunkMarkdownFile(parsed.file, {
58    maxWords: parsed.maxWords,
59    outputDir: parsed.outputDir,
60  })
61 
62  console.log(JSON.stringify(result))
63  return 0
64}
65 
66export function chunkMarkdownFile(
67  file: string,
68  options: { maxWords?: number; outputDir?: string } = {}
69): ChunkResult {
70  const maxWords = options.maxWords ?? 5000
71  const outputDir = options.outputDir ?? ""
72 
73  const rawContent = normalizeNewlines(readFileSync(file, "utf-8"))
74  const { frontmatter, body } = extractFrontmatter(rawContent)
75  const chunks = buildChunks(parseMarkdown(body), maxWords)
76 
77  const dir = outputDir ? join(outputDir, "chunks") : join(dirname(file), "chunks")
78  mkdirSync(dir, { recursive: true })
79 
80  if (frontmatter) {
81    writeFileSync(join(dir, "frontmatter.md"), frontmatter)
82  }
83 
84  chunks.forEach((chunk, index) => {
85    const num = String(index + 1).padStart(2, "0")
86    writeFileSync(join(dir, `chunk-${num}.md`), chunk.blocks.map(block => block.md).join("\n\n"))
87  })
88 
89  return {
90    source: file,
91    chunks: chunks.length,
92    output_dir: dir,
93    frontmatter: Boolean(frontmatter),
94    words_per_chunk: chunks.map(chunk => chunk.words),
95  }
96}
97 
98function parseChunkCliArgs(args: string[]):
99  | ChunkCliOptions
100  | { help: true }
101  | { error: string } {
102  let file = ""
103  let maxWords = 5000
104  let outputDir = ""
105 
106  for (let index = 0; index < args.length; index += 1) {
107    const arg = args[index]
108 
109    if (arg === "-h" || arg === "--help") {
110      return { help: true }
111    }
112 
113    if (arg === "--max-words") {
114      const value = args[index + 1]
115      if (!value) return { error: "Missing value for --max-words" }
116      maxWords = parsePositiveInt(value, 0)
117      if (maxWords <= 0) return { error: `Invalid --max-words value: ${value}` }
118      index += 1
119      continue
120    }
121 
122    if (arg === "--output-dir") {
123      const value = args[index + 1]
124      if (!value) return { error: "Missing value for --output-dir" }
125      outputDir = value
126      index += 1
127      continue
128    }
129 
130    if (arg.startsWith("-")) {
131      return { error: `Unknown option: ${arg}` }
132    }
133 
134    if (!file) {
135      file = arg
136      continue
137    }
138 
139    return { error: `Unexpected positional argument: ${arg}` }
140  }
141 
142  if (!file) {
143    return { error: "Missing input file" }
144  }
145 
146  return { file, maxWords, outputDir }
147}
148 
149function parsePositiveInt(value: string | undefined, fallback: number): number {
150  if (!value) return fallback
151  const parsed = Number.parseInt(value, 10)
152  return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback
153}
154 
155function normalizeNewlines(text: string): string {
156  return text.replace(/^\uFEFF/, "").replace(/\r\n?/g, "\n")
157}
158 
159function trimBoundaryBlankLines(text: string): string {
160  return text.replace(/^\n+/, "").replace(/\n+$/, "")
161}
162 
163function extractFrontmatter(content: string): { frontmatter: string; body: string } {
164  const lines = content.split("\n")
165  if (lines[0] !== "---") {
166    return { frontmatter: "", body: content }
167  }
168 
169  for (let index = 1; index < lines.length; index += 1) {
170    if (lines[index] === "---" || lines[index] === "...") {
171      return {
172        frontmatter: lines.slice(0, index + 1).join("\n"),
173        body: lines.slice(index + 1).join("\n").replace(/^\n+/, ""),
174      }
175    }
176  }
177 
178  return { frontmatter: "", body: content }
179}
180 
181function parseMarkdown(content: string): Block[] {
182  if (!content.trim()) return []
183 
184  const lines = content.split("\n")
185  const tokens = parser.parse(content, {})
186  const blocks: Block[] = []
187 
188  for (const token of tokens) {
189    if (!token.map || token.level !== 0) continue
190    if (token.nesting !== 1 && token.nesting !== 0) continue
191 
192    const [startLine, endLine] = token.map
193    const md = trimBoundaryBlankLines(lines.slice(startLine, endLine).join("\n"))
194    if (!md) continue
195 
196    blocks.push(makeBlock(tokenTypeToBlockKind(token.type), md))
197  }
198 
199  if (blocks.length === 0) {
200    const body = trimBoundaryBlankLines(content)
201    if (body) {
202      blocks.push(makeBlock("flow", body))
203    }
204  }
205 
206  return blocks
207}
208 
209function tokenTypeToBlockKind(tokenType: string): BlockKind {
210  if (tokenType === "heading_open") return "heading"
211  if (tokenType === "hr") return "thematicBreak"
212  if (tokenType === "html_block") return "html"
213  if (tokenType === "fence" || tokenType === "code_block") return "code"
214  return "flow"
215}
216 
217function makeBlock(kind: BlockKind, md: string): Block {
218  return {
219    kind,
220    md: trimBoundaryBlankLines(md),
221    words: countWords(md),
222  }
223}
224 
225function buildChunks(blocks: Block[], maxWordsPerChunk: number): Chunk[] {
226  const sections = splitIntoSections(blocks)
227  const normalizedBlocks: Block[] = []
228 
229  for (const section of sections) {
230    const sectionWords = section.reduce((sum, block) => sum + block.words, 0)
231    if (sectionWords <= maxWordsPerChunk) {
232      normalizedBlocks.push(makeBlock("flow", section.map(block => block.md).join("\n\n")))
233      continue
234    }
235 
236    for (const block of section) {
237      normalizedBlocks.push(...splitOversizedBlock(block, maxWordsPerChunk))
238    }
239  }
240 
241  const chunks: Chunk[] = []
242  let currentBlocks: Block[] = []
243  let currentWords = 0
244 
245  for (const block of normalizedBlocks) {
246    if (currentWords + block.words > maxWordsPerChunk && currentBlocks.length > 0) {
247      chunks.push({ blocks: currentBlocks, words: currentWords })
248      currentBlocks = [block]
249      currentWords = block.words
250      continue
251    }
252 
253    currentBlocks.push(block)
254    currentWords += block.words
255  }
256 
257  if (currentBlocks.length > 0) {
258    chunks.push({ blocks: currentBlocks, words: currentWords })
259  }
260 
261  return chunks
262}
263 
264function splitIntoSections(blocks: Block[]): Block[][] {
265  const sections: Block[][] = []
266  let current: Block[] = []
267 
268  for (const block of blocks) {
269    if (block.kind === "heading" && current.length > 0) {
270      sections.push(current)
271      current = [block]
272      continue
273    }
274 
275    current.push(block)
276  }
277 
278  if (current.length > 0) {
279    sections.push(current)
280  }
281 
282  return sections
283}
284 
285function splitOversizedBlock(block: Block, maxWordsPerChunk: number): Block[] {
286  if (block.words <= maxWordsPerChunk) return [block]
287 
288  if (
289    block.kind === "heading"
290    || block.kind === "thematicBreak"
291    || block.kind === "html"
292    || block.kind === "code"
293  ) {
294    return [block]
295  }
296 
297  const lines = block.md.split("\n")
298  if (lines.length <= 1) {
299    return [block]
300  }
301 
302  const splitBlocks: Block[] = []
303  let buffer: string[] = []
304  let bufferWords = 0
305 
306  for (const line of lines) {
307    const lineWords = countWords(line)
308    if (bufferWords + lineWords > maxWordsPerChunk && buffer.length > 0) {
309      splitBlocks.push(makeBlock(block.kind, buffer.join("\n")))
310      buffer = [line]
311      bufferWords = lineWords
312      continue
313    }
314 
315    buffer.push(line)
316    bufferWords += lineWords
317  }
318 
319  if (buffer.length > 0) {
320    splitBlocks.push(makeBlock(block.kind, buffer.join("\n")))
321  }
322 
323  return splitBlocks
324}
325 
326function countWords(text: string): number {
327  const cleaned = text.replace(/[#*`\[\]()>|_~-]/g, " ")
328  const cjk = cleaned.match(/[\u4e00-\u9fff\u3400-\u4dbf\uf900-\ufaff]/g)
329  const latin = cleaned.match(/[a-zA-Z0-9]+/g)
330  return (cjk?.length ?? 0) + (latin?.length ?? 0)
331}
332 
333if (import.meta.main) {
334  process.exit(runChunkCli(process.argv.slice(2), process.argv[1] ?? "chunk.ts"))
335}
336