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

package websearch

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

	"img_infinite_canvas/internal/domain/design"
)

type DuckDuckGoOptions struct {
	TimeoutSeconds int
	MaxResults     int
}

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

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
	}
	return &DuckDuckGoResearcher{
		client:     &http.Client{Timeout: timeout},
		maxResults: maxResults,
	}
}

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")
	}
	return report, 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 = html.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 = html.UnescapeString(value)
	value = spaceRe.ReplaceAllString(value, " ")
	return strings.TrimSpace(value)
}

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, "/")
}