moteva/server/internal/infrastructure/websearch/duckduckgo.go

package websearch

import (
	"bytes"
	"context"
	"fmt"
	stdhtml "html"
	"io"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"

	"img_infinite_canvas/internal/domain/design"

	xhtml "golang.org/x/net/html"
)

type DuckDuckGoOptions struct {
	TimeoutSeconds      int
	MaxResults          int
	MaxPageContentRunes int
}

type DuckDuckGoResearcher struct {
	client              *http.Client
	maxResults          int
	maxPageContentRunes int
}

const (
	defaultMaxPageContentRunes = 1800
	maxPageContentBytes        = 2 << 20
	maxPageContentConcurrency  = 4
)

func NewDuckDuckGoResearcher(opts DuckDuckGoOptions) *DuckDuckGoResearcher {
	timeout := time.Duration(opts.TimeoutSeconds) * time.Second
	if timeout <= 0 {
		timeout = 10 * time.Second
	}
	maxResults := opts.MaxResults
	if maxResults <= 0 {
		maxResults = 5
	}
	maxPageContentRunes := opts.MaxPageContentRunes
	if maxPageContentRunes <= 0 {
		maxPageContentRunes = defaultMaxPageContentRunes
	}
	return &DuckDuckGoResearcher{
		client:              &http.Client{Timeout: timeout},
		maxResults:          maxResults,
		maxPageContentRunes: maxPageContentRunes,
	}
}

func (r *DuckDuckGoResearcher) Research(ctx context.Context, prompt string) (design.ResearchReport, error) {
	query := BuildVisualSearchQuery(prompt)
	report := design.ResearchReport{
		Kind:   "web_search",
		Source: "duckduckgo",
		Query:  query,
	}

	endpoint := "https://html.duckduckgo.com/html/?q=" + url.QueryEscape(query)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
	if err != nil {
		return report, err
	}
	req.Header.Set("User-Agent", "Mozilla/5.0 (compatible; MotevaDesignAgent/1.0)")
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8")

	resp, err := r.client.Do(req)
	if err != nil {
		return report, err
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return report, fmt.Errorf("web search failed: %s", resp.Status)
	}

	body, err := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
	if err != nil {
		return report, err
	}

	report.Results = parseDuckDuckGoResults(string(body), r.maxResults)
	if len(report.Results) == 0 {
		return report, fmt.Errorf("web search returned no results")
	}
	r.hydrateResultContents(ctx, report.Results)
	return report, nil
}

func (r *DuckDuckGoResearcher) hydrateResultContents(ctx context.Context, results []design.ResearchResult) {
	if r.maxPageContentRunes <= 0 || len(results) == 0 {
		return
	}

	var wg sync.WaitGroup
	sem := make(chan struct{}, maxPageContentConcurrency)
	for index := range results {
		if strings.TrimSpace(results[index].URL) == "" {
			continue
		}
		wg.Add(1)
		go func(index int) {
			defer wg.Done()
			select {
			case sem <- struct{}{}:
				defer func() { <-sem }()
			case <-ctx.Done():
				return
			}
			content, err := r.fetchPageContent(ctx, results[index].URL)
			if err == nil && content != "" {
				results[index].Content = content
			}
		}(index)
	}
	wg.Wait()
}

func (r *DuckDuckGoResearcher) fetchPageContent(ctx context.Context, rawURL string) (string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
	if err != nil {
		return "", err
	}
	req.Header.Set("User-Agent", "Mozilla/5.0 (compatible; MotevaDesignAgent/1.0)")
	req.Header.Set("Accept", "text/html,application/xhtml+xml,text/plain;q=0.9,*/*;q=0.3")

	resp, err := r.client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return "", fmt.Errorf("fetch page failed: %s", resp.Status)
	}

	contentType := strings.ToLower(resp.Header.Get("Content-Type"))
	if contentType != "" &&
		!strings.Contains(contentType, "text/html") &&
		!strings.Contains(contentType, "application/xhtml+xml") &&
		!strings.Contains(contentType, "text/plain") {
		return "", nil
	}

	body, err := io.ReadAll(io.LimitReader(resp.Body, maxPageContentBytes))
	if err != nil {
		return "", err
	}
	if strings.Contains(contentType, "text/plain") {
		return cleanPageContentText(string(body), r.maxPageContentRunes), nil
	}
	return extractReadableHTMLText(body, r.maxPageContentRunes), nil
}

func BuildVisualSearchQuery(prompt string) string {
	prompt = cleanPromptForSearch(prompt)
	if prompt == "" {
		return "brand visual design website inspiration"
	}

	lower := strings.ToLower(prompt)
	if isLogoReferencePrompt(prompt, lower) {
		return buildLogoReferenceQuery(prompt, lower)
	}

	keywords := []string{"website design reference", "brand visual style", "ecommerce ui"}
	if strings.Contains(lower, "hoodie") || strings.Contains(prompt, "服装") || strings.Contains(prompt, "针织") {
		keywords = []string{"fashion ecommerce website design", "hoodie brand visual reference", "streetwear UI"}
	}
	if strings.Contains(prompt, "海报") || strings.Contains(lower, "poster") {
		keywords = []string{"poster design reference", "campaign key visual", "layout inspiration"}
	}
	if strings.Contains(prompt, "包装") || strings.Contains(lower, "packaging") {
		keywords = []string{"packaging design reference", "brand identity system", "visual guidelines"}
	}
	if strings.Contains(lower, "ui") || strings.Contains(prompt, "界面") {
		keywords = []string{"UI design reference", "product interface inspiration", "case study"}
	}

	return strings.Join(append([]string{prompt}, keywords...), " ")
}

func isLogoReferencePrompt(prompt string, lower string) bool {
	return strings.Contains(lower, "logo") ||
		strings.Contains(prompt, "品牌标识") ||
		strings.Contains(prompt, "标志") ||
		strings.Contains(prompt, "商标")
}

func buildLogoReferenceQuery(prompt string, lower string) string {
	subject := logoReferenceSubject(prompt, lower)
	terms := []string{
		"logo design case study",
		"brand identity",
		"visual identity",
		"design process",
		"moodboard inspiration",
		"Behance Dribbble",
		"-maker",
		"-generator",
		"-free",
		"-online",
		"-template",
		"-工具",
		"-生成器",
		"-制作器",
		"-模板",
		"-免费",
	}
	if hasCJK(prompt) {
		terms = append([]string{
			"logo 设计稿",
			"品牌视觉",
			"设计思路",
			"案例",
			"灵感",
		}, terms...)
	}
	if subject != "" {
		terms = append([]string{subject}, terms...)
	}
	return strings.Join(terms, " ")
}

func logoReferenceSubject(prompt string, lower string) string {
	if strings.Contains(lower, "design agent") ||
		strings.Contains(lower, "design-agent") ||
		strings.Contains(prompt, "设计 Agent") ||
		strings.Contains(prompt, "设计Agent") ||
		strings.Contains(prompt, "设计智能体") ||
		strings.Contains(prompt, "设计助手") {
		return "AI design agent"
	}
	if strings.Contains(lower, "ai") && (strings.Contains(lower, "design") || strings.Contains(prompt, "设计")) {
		return "AI design product"
	}

	subject := prompt
	replacer := strings.NewReplacer(
		"制作一个", " ",
		"制作一款", " ",
		"设计一个", " ",
		"设计一款", " ",
		"做一个", " ",
		"做一款", " ",
		"生成一个", " ",
		"生成一款", " ",
		"创建一个", " ",
		"创建一款", " ",
		"logo", " ",
		"Logo", " ",
		"LOGO", " ",
		"品牌标识", " ",
		"标志", " ",
		"商标", " ",
		"让设计变简单", " ",
		"帮你搞定一切设计需求", " ",
		"一切设计需求", " ",
		"create a", " ",
		"make a", " ",
		"design a", " ",
		"generate a", " ",
	)
	subject = replacer.Replace(subject)
	subject = strings.Trim(subject, " ,，。.!！?？;；:：-—_")
	fields := strings.Fields(subject)
	if len(fields) > 8 {
		fields = fields[:8]
	}
	return strings.Join(fields, " ")
}

func hasCJK(value string) bool {
	for _, r := range value {
		if r >= '\u4e00' && r <= '\u9fff' {
			return true
		}
	}
	return false
}

func cleanPromptForSearch(prompt string) string {
	lines := strings.Split(prompt, "\n")
	cleaned := make([]string, 0, len(lines))
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		if strings.HasPrefix(line, "模型选择") || strings.HasPrefix(line, "Selected model:") {
			continue
		}
		if strings.HasPrefix(line, "请先进行联网参考") || strings.HasPrefix(line, "Start with web references") {
			continue
		}
		if strings.HasPrefix(line, "参考图") || strings.HasPrefix(line, "Reference image") {
			continue
		}
		cleaned = append(cleaned, line)
	}
	value := strings.Join(cleaned, " ")
	fields := strings.Fields(value)
	if len(fields) > 40 {
		fields = fields[:40]
	}
	return strings.Join(fields, " ")
}

var (
	resultLinkRe = regexp.MustCompile(`(?is)<a[^>]*class="[^"]*\bresult__a\b[^"]*"[^>]*href="([^"]+)"[^>]*>(.*?)</a>`)
	snippetRe    = regexp.MustCompile(`(?is)<a[^>]*class="[^"]*\bresult__snippet\b[^"]*"[^>]*>(.*?)</a>`)
	tagRe        = regexp.MustCompile(`(?is)<[^>]+>`)
	spaceRe      = regexp.MustCompile(`\s+`)
)

func parseDuckDuckGoResults(body string, maxResults int) []design.ResearchResult {
	matches := resultLinkRe.FindAllStringSubmatchIndex(body, -1)
	results := make([]design.ResearchResult, 0, maxResults)
	seen := map[string]bool{}

	for index, match := range matches {
		if len(results) >= maxResults {
			break
		}

		rawURL := body[match[2]:match[3]]
		titleHTML := body[match[4]:match[5]]
		resolvedURL := resolveDuckDuckGoURL(rawURL)
		if resolvedURL == "" {
			continue
		}

		key := dedupeKey(resolvedURL)
		if key == "" || seen[key] {
			continue
		}

		nextStart := len(body)
		if index+1 < len(matches) {
			nextStart = matches[index+1][0]
		}
		segment := body[match[1]:nextStart]
		snippet := ""
		if snippetMatch := snippetRe.FindStringSubmatch(segment); len(snippetMatch) > 1 {
			snippet = cleanHTML(snippetMatch[1])
		}

		title := cleanHTML(titleHTML)
		if title == "" {
			continue
		}
		seen[key] = true
		results = append(results, design.ResearchResult{
			Title:   title,
			URL:     resolvedURL,
			Snippet: snippet,
		})
	}

	return results
}

func resolveDuckDuckGoURL(raw string) string {
	raw = stdhtml.UnescapeString(strings.TrimSpace(raw))
	if strings.HasPrefix(raw, "//") {
		raw = "https:" + raw
	}

	parsed, err := url.Parse(raw)
	if err != nil {
		return ""
	}
	if strings.Contains(parsed.Host, "duckduckgo.com") {
		if uddg := parsed.Query().Get("uddg"); uddg != "" {
			raw = uddg
		}
	}
	if !strings.HasPrefix(raw, "http://") && !strings.HasPrefix(raw, "https://") {
		return ""
	}
	return raw
}

func cleanHTML(value string) string {
	value = tagRe.ReplaceAllString(value, " ")
	value = stdhtml.UnescapeString(value)
	value = spaceRe.ReplaceAllString(value, " ")
	return strings.TrimSpace(value)
}

func extractReadableHTMLText(body []byte, maxRunes int) string {
	doc, err := xhtml.Parse(bytes.NewReader(body))
	if err != nil {
		return cleanPageContentText(cleanHTML(string(body)), maxRunes)
	}

	parts := make([]string, 0, 64)
	runeCount := 0
	collectReadableText(doc, false, &parts, &runeCount, maxRunes)
	return cleanPageContentText(strings.Join(parts, " "), maxRunes)
}

func collectReadableText(node *xhtml.Node, skip bool, parts *[]string, runeCount *int, maxRunes int) bool {
	if node == nil || *runeCount >= maxRunes {
		return true
	}
	if node.Type == xhtml.ElementNode && skippedPageTextTag(strings.ToLower(node.Data)) {
		skip = true
	}
	if !skip && node.Type == xhtml.TextNode {
		text := cleanPageContentText(node.Data, maxRunes)
		if usefulPageText(text) {
			*parts = append(*parts, text)
			*runeCount += len([]rune(text)) + 1
			if *runeCount >= maxRunes {
				return true
			}
		}
	}
	for child := node.FirstChild; child != nil; child = child.NextSibling {
		if collectReadableText(child, skip, parts, runeCount, maxRunes) {
			return true
		}
	}
	return false
}

func skippedPageTextTag(tag string) bool {
	switch tag {
	case "head", "script", "style", "noscript", "template", "svg", "canvas", "iframe", "form", "button", "nav", "header", "footer":
		return true
	default:
		return false
	}
}

func usefulPageText(text string) bool {
	text = strings.TrimSpace(text)
	if len([]rune(text)) < 2 {
		return false
	}
	lower := strings.ToLower(text)
	if len([]rune(text)) < 80 && (strings.Contains(lower, "cookie") || strings.Contains(lower, "privacy policy")) {
		return false
	}
	return true
}

func cleanPageContentText(value string, maxRunes int) string {
	value = stdhtml.UnescapeString(value)
	value = spaceRe.ReplaceAllString(value, " ")
	value = strings.TrimSpace(value)
	return clampRunes(value, maxRunes)
}

func clampRunes(value string, maxRunes int) string {
	if maxRunes <= 0 {
		return value
	}
	runes := []rune(value)
	if len(runes) <= maxRunes {
		return value
	}
	return strings.TrimSpace(string(runes[:maxRunes]))
}

func dedupeKey(raw string) string {
	parsed, err := url.Parse(raw)
	if err != nil {
		return ""
	}
	host := strings.TrimPrefix(strings.ToLower(parsed.Host), "www.")
	return host + strings.TrimRight(parsed.Path, "/")
}