package ai

import (
	"regexp"
	"strings"
	"unicode/utf8"
)

var (
	anyFence     = regexp.MustCompile("(?s)```[a-zA-Z]*\\s*\n.*?```")
	cotCue       = regexp.MustCompile(`(?i)(我们需要|作为助手|指令要求|让我们看看|最终答案|chain of thought|step by step|I need to|as an assistant|the instruction)`)
	sqlFenceFind = regexp.MustCompile("(?s)```sql\\s*\n(.+?)```")
)

// SanitizeAssistantContent 去掉 Ask/Generate 回复里的推理草稿，只保留短说明 + SQL 块（或短澄清）。
func SanitizeAssistantContent(content string) string {
	s := strings.TrimSpace(content)
	if s == "" {
		return s
	}
	fences := anyFence.FindAllStringIndex(s, -1)
	if len(fences) == 0 {
		if utf8.RuneCountInString(s) <= 240 && !cotCue.MatchString(s) {
			return s
		}
		return keepLeadingLines(s, 2)
	}
	var b strings.Builder
	prev := 0
	for i, loc := range fences {
		prose := strings.TrimSpace(s[prev:loc[0]])
		if prose != "" {
			if i == 0 {
				prose = trimProse(prose)
			} else {
				prose = keepLeadingLines(prose, 2)
			}
			if prose != "" {
				if b.Len() > 0 {
					b.WriteByte('\n')
				}
				b.WriteString(prose)
			}
		}
		if b.Len() > 0 {
			b.WriteString("\n\n")
		}
		b.WriteString(strings.TrimSpace(s[loc[0]:loc[1]]))
		prev = loc[1]
	}
	tail := strings.TrimSpace(s[prev:])
	if tail != "" {
		tail = keepLeadingLines(tail, 2)
		if tail != "" {
			b.WriteString("\n\n")
			b.WriteString(tail)
		}
	}
	out := strings.TrimSpace(b.String())
	if out == "" {
		if m := sqlFenceFind.FindString(s); m != "" {
			return strings.TrimSpace(m)
		}
		return keepLeadingLines(s, 2)
	}
	return out
}

func trimProse(prose string) string {
	lines := nonEmptyLines(prose)
	if len(lines) == 0 {
		return ""
	}
	var notes []string
	for _, ln := range lines {
		if strings.HasPrefix(ln, "说明") || strings.HasPrefix(ln, "注意") || strings.HasPrefix(ln, "提示") {
			notes = append(notes, ln)
		}
	}
	if len(notes) > 0 {
		if len(notes) > 2 {
			notes = notes[len(notes)-2:]
		}
		return strings.Join(notes, "\n")
	}
	joined := strings.Join(lines, "\n")
	if cotCue.MatchString(joined) || utf8.RuneCountInString(joined) > 160 {
		return ""
	}
	return keepLeadingLines(joined, 2)
}

func keepLeadingLines(s string, n int) string {
	lines := nonEmptyLines(s)
	if len(lines) > n {
		lines = lines[:n]
	}
	return strings.Join(lines, "\n")
}

func nonEmptyLines(s string) []string {
	raw := strings.Split(s, "\n")
	out := make([]string, 0, len(raw))
	for _, ln := range raw {
		ln = strings.TrimSpace(ln)
		if ln != "" {
			out = append(out, ln)
		}
	}
	return out
}