package rag

import (
	"strings"
)

// 切块参数（针对中文 + embedding token 上限调过；后续可配置化）。
const (
	chunkTargetRunes  = 500  // 目标大小：贪心打包到接近此值就在句界收口
	chunkOverlapRunes = 80   // 块间重叠：保跨块上下文连续
	chunkMinRunes     = 100  // 最小块：低于则并入相邻，避免碎块稀释向量
	chunkMaxRunes     = 1000 // 原子硬上限：超大无标点段落兜底窗口切
)

// chunk 把文本切成检索友好的语义块。Markdown 文档走「标题感知」切块：
// 按 #/## 层级切段，块不跨章节，且每块前缀完整标题路径（如「部署 > 环境要求」）补全上下文；
// 非 Markdown 走纯语义切块。全程按 rune 操作，杜绝中文 UTF-8 被字节切碎。
func chunk(text string) []string {
	if hasMarkdownHeadings(text) {
		return chunkMarkdown(text)
	}
	return packSection(text)
}

// packSection 对一段纯文本做语义切块：原子 → 贪心打包（句末收口）→ 块间重叠。
func packSection(text string) []string {
	return addOverlap(packAtoms(splitToAtoms(text), chunkTargetRunes, chunkMinRunes), chunkOverlapRunes)
}

// section 是一个标题层级下的正文段：crumb 为标题路径面包屑（顶到本节，可空=前言）。
type section struct {
	crumb string
	body  string
}

// chunkMarkdown 标题感知切块：按标题层级切成 section（块绝不跨章节），每节内部再语义切块，
// 并给每块前缀该节的标题路径——使「## Redis 下的『端口配置』」这类块自带章节语境，显著提升检索命中。
func chunkMarkdown(text string) []string {
	var out []string
	for _, s := range splitMarkdownSections(text) {
		for _, body := range packSection(s.body) {
			if s.crumb != "" {
				out = append(out, s.crumb+"\n"+body) // 面包屑作为上下文前缀
			} else {
				out = append(out, body)
			}
		}
	}
	return out
}

// splitMarkdownSections 按 ATX 标题（# ~ ######）把文档切成带面包屑的章节，正确跳过代码围栏内的 #。
func splitMarkdownSections(text string) []section {
	lines := strings.Split(strings.ReplaceAll(text, "\r\n", "\n"), "\n")
	type heading struct {
		level int
		title string
	}
	var stack []heading
	var sections []section
	var body []string
	inFence, fenceMarker := false, ""

	crumb := func() string {
		parts := make([]string, len(stack))
		for i, h := range stack {
			parts[i] = h.title
		}
		return strings.Join(parts, " > ")
	}
	flush := func() {
		if b := strings.TrimSpace(strings.Join(body, "\n")); b != "" {
			sections = append(sections, section{crumb: crumb(), body: b})
		}
		body = body[:0]
	}

	for _, ln := range lines {
		t := strings.TrimSpace(ln)
		if m := fenceMark(t); m != "" { // 代码围栏开/合：围栏内的 # 不当标题
			if !inFence {
				inFence, fenceMarker = true, m
			} else if m == fenceMarker {
				inFence, fenceMarker = false, ""
			}
			body = append(body, ln)
			continue
		}
		if !inFence {
			if lvl, title, ok := parseHeading(t); ok {
				flush() // 上一节正文收口
				for len(stack) > 0 && stack[len(stack)-1].level >= lvl {
					stack = stack[:len(stack)-1] // 同级/更浅的标题出栈
				}
				stack = append(stack, heading{lvl, title})
				continue
			}
		}
		body = append(body, ln)
	}
	flush()
	return sections
}

// parseHeading 解析 ATX 标题行：返回层级(1-6)、标题文本、是否命中。要求 # 后跟空格（排除 #define 之类）。
func parseHeading(line string) (level int, title string, ok bool) {
	n := 0
	for n < len(line) && line[n] == '#' {
		n++
	}
	if n == 0 || n > 6 || n >= len(line) || line[n] != ' ' {
		return 0, "", false
	}
	title = strings.TrimRight(strings.TrimSpace(line[n:]), "# ") // 去掉闭合式 ATX 尾部的 #
	if title == "" {
		return 0, "", false
	}
	return n, title, true
}

// fenceMark 返回代码围栏标记（``` 或 ~~~），非围栏行返回空。
func fenceMark(t string) string {
	switch {
	case strings.HasPrefix(t, "```"):
		return "```"
	case strings.HasPrefix(t, "~~~"):
		return "~~~"
	default:
		return ""
	}
}

// hasMarkdownHeadings 快速判断文本是否含 ATX 标题（决定走标题感知还是纯语义切块）。
func hasMarkdownHeadings(text string) bool {
	for _, ln := range strings.Split(text, "\n") {
		if _, _, ok := parseHeading(strings.TrimSpace(ln)); ok {
			return true
		}
	}
	return false
}

// splitToAtoms 把文本切成"原子"（句子/行）：在换行与句末标点处断开，去空白；
// 超大无标点原子按 rune 窗口兜底切到 ≤ 目标大小。原子是打包的最小不可分单元。
func splitToAtoms(text string) []string {
	runes := []rune(strings.ReplaceAll(text, "\r\n", "\n"))
	var atoms []string
	start := 0
	flush := func(end int) {
		if s := strings.TrimSpace(string(runes[start:end])); s != "" {
			atoms = append(atoms, s)
		}
		start = end
	}
	for i := 0; i < len(runes); i++ {
		r := runes[i]
		switch {
		case r == '\n' || isCJKEnd(r):
			flush(i + 1)
		case r == '.' || r == '!' || r == '?' || r == ';':
			// ASCII 句末：仅当其后为空白/行尾才断（避开缩写、小数点）。
			if i+1 >= len(runes) || runes[i+1] == ' ' || runes[i+1] == '\n' {
				flush(i + 1)
			}
		}
	}
	flush(len(runes))

	out := make([]string, 0, len(atoms))
	for _, a := range atoms {
		if runeLen(a) <= chunkMaxRunes {
			out = append(out, a)
		} else {
			out = append(out, splitByRuneWindow(a, chunkTargetRunes)...)
		}
	}
	return out
}

// packAtoms 贪心打包：把原子拼进当前块，直到再加会超 target 就收口（块在句界结束）。
// 末尾过小的块（< min）并入前一块，避免碎块。
func packAtoms(atoms []string, target, min int) []string {
	var chunks []string
	var cur strings.Builder
	curLen := 0
	closeCur := func() {
		if curLen > 0 {
			chunks = append(chunks, cur.String())
			cur.Reset()
			curLen = 0
		}
	}
	for _, a := range atoms {
		al := runeLen(a)
		if curLen > 0 && curLen+al > target {
			closeCur()
		}
		if curLen > 0 {
			cur.WriteByte('\n')
			curLen++
		}
		cur.WriteString(a)
		curLen += al
	}
	if curLen > 0 {
		if n := len(chunks); n > 0 && curLen < min {
			chunks[n-1] = chunks[n-1] + "\n" + cur.String() // 尾块太小 → 并入上一块
		} else {
			chunks = append(chunks, cur.String())
		}
	}
	return chunks
}

// addOverlap 给每块前缀上一块尾部的 overlap 个 rune，保跨块上下文连续（首块不加）。
func addOverlap(chunks []string, overlap int) []string {
	if overlap <= 0 || len(chunks) <= 1 {
		return chunks
	}
	out := make([]string, len(chunks))
	out[0] = chunks[0]
	for i := 1; i < len(chunks); i++ {
		prev := []rune(chunks[i-1])
		tail := prev
		if len(prev) > overlap {
			tail = prev[len(prev)-overlap:]
		}
		out[i] = strings.TrimSpace(string(tail)) + "\n" + chunks[i]
	}
	return out
}

func isCJKEnd(r rune) bool {
	switch r {
	case '。', '！', '？', '；', '…':
		return true
	}
	return false
}

// splitByRuneWindow 按 rune 窗口硬切（兜底：超大无标点原子），保证不切碎多字节字符。
func splitByRuneWindow(s string, size int) []string {
	r := []rune(s)
	var out []string
	for len(r) > size {
		if t := strings.TrimSpace(string(r[:size])); t != "" {
			out = append(out, t)
		}
		r = r[size:]
	}
	if t := strings.TrimSpace(string(r)); t != "" {
		out = append(out, t)
	}
	return out
}

func runeLen(s string) int { return len([]rune(s)) }