geoGo/internal/collect/doubao.go

package collect

import (
	"context"
	"fmt"
	"geo/internal/config"
	"log"
	"strings"
	"time"

	"github.com/go-rod/rod"
	"github.com/go-rod/rod/lib/proto"
)

// DoubaoCollector 豆包收集器
type DoubaoCollector struct {
	*BaseCollector
}

// NewDoubaoCollector 创建豆包收集器
func NewDoubaoCollector(ctx context.Context, params *CollectParams, cfg *config.Config, logger *log.Logger) CollectorInterface {
	collector := &DoubaoCollector{
		BaseCollector: NewBaseCollector(ctx, params, cfg, logger),
	}

	// 设置豆包的URL
	collector.LoginURL = "https://www.doubao.com/"
	collector.ChatURL = "https://www.doubao.com/chat/"

	return collector
}

// CheckLoginStatus 检查登录状态
func (c *DoubaoCollector) CheckLoginStatus() bool {
	currentURL := c.GetCurrentURL()

	// 检查是否在聊天页面
	if strings.Contains(currentURL, "doubao.com") {
		// 查找用户信息元素
		userInfo, err := c.SafeElement(".user-info, .avatar, [class*='user-profile']")
		if err == nil && userInfo != nil {
			return true
		}

		// 检查是否有输入框
		inputBox, err := c.SafeElement("textarea, [contenteditable='true']")
		if err == nil && inputBox != nil {
			return true
		}
	}

	return false
}

// WaitLogin 等待登录
func (c *DoubaoCollector) WaitLogin() (bool, string) {
	c.LogInfo("开始等待豆包登录...")

	if err := c.SetupDriver(); err != nil {
		return false, fmt.Sprintf("浏览器启动失败: %v", err)
	}
	defer c.Close()

	// 访问豆包首页
	c.Page.MustNavigate(c.LoginURL)
	c.Sleep(3)

	// 检查是否已登录
	if c.CheckLoginStatus() {
		c.SaveCookies()
		c.LogInfo("已有登录状态")
		return true, "already_logged_in"
	}

	c.LogInfo("请登录豆包账号...")

	// 等待用户手动登录，最多300秒
	for i := 0; i < 300; i++ {
		if c.CheckLoginStatus() {
			c.SaveCookies()
			c.LogInfo("登录成功")
			return true, "login_success"
		}

		time.Sleep(1 * time.Second)
	}

	return false, "登录超时，请检查网络或账号状态"
}

// AskQuestion 提问并获取答案
func (c *DoubaoCollector) AskQuestion(question string) (string, error) {
	c.LogInfo(fmt.Sprintf("开始向豆包提问: %s", question))

	// 初始化浏览器
	if err := c.SetupDriver(); err != nil {
		return "", fmt.Errorf("浏览器启动失败: %v", err)
	}
	defer c.Close()

	// 初始化页面
	if err := c.InitPage(); err != nil {
		return "", fmt.Errorf("页面初始化失败，请先调用WaitLogin登录: %v", err)
	}

	c.Sleep(3)

	// 输入问题
	if err := c.inputQuestion(question); err != nil {
		return "", fmt.Errorf("输入问题失败: %v", err)
	}

	// 点击发送
	if err := c.clickSendButton(); err != nil {
		return "", fmt.Errorf("点击发送按钮失败: %v", err)
	}

	// 等待并获取答案
	answer, err := c.waitForAnswer()
	if err != nil {
		return "", fmt.Errorf("获取答案失败: %v", err)
	}

	c.LogInfo(fmt.Sprintf("成功获取豆包答案，长度: %d 字符", len(answer)))
	return answer, nil
}

// inputQuestion 输入问题
func (c *DoubaoCollector) inputQuestion(question string) error {
	c.LogInfo("输入问题到豆包...")

	// 豆包的输入框选择器
	inputSelectors := []string{
		"textarea[placeholder*='输入']",
		"textarea[placeholder*='问']",
		"textarea",
		"[contenteditable='true']",
		".chat-input textarea",
		"#input-box",
		".input-area textarea",
	}

	var inputBox *rod.Element
	var err error

	for _, selector := range inputSelectors {
		inputBox, err = c.WaitForElementVisible(selector, 10)
		if err == nil && inputBox != nil {
			c.LogInfo(fmt.Sprintf("找到输入框: %s", selector))
			break
		}
	}

	if inputBox == nil {
		return fmt.Errorf("未找到输入框")
	}

	// 点击获取焦点
	if err := inputBox.Click(proto.InputMouseButtonLeft, 1); err != nil {
		return fmt.Errorf("点击输入框失败: %v", err)
	}
	c.SleepMs(500)

	// 清空输入框
	if err := c.ClearInput(inputBox); err != nil {
		c.LogInfo(fmt.Sprintf("清空输入框失败: %v", err))
	}
	c.SleepMs(300)

	// 输入问题
	if err := c.SetInputValue(inputBox, question); err != nil {
		inputBox.Input(question)
	}

	c.LogInfo(fmt.Sprintf("问题已输入"))
	c.SleepMs(1000)

	return nil
}

// clickSendButton 点击发送按钮
func (c *DoubaoCollector) clickSendButton() error {
	c.LogInfo("点击发送按钮...")

	// 发送按钮选择器
	sendSelectors := []string{
		"button[class*='send']",
		"button[class*='submit']",
		".send-btn",
		".submit-btn",
		"button svg[path*='send']",
		"[aria-label*='发送']",
		".send-icon",
	}

	var sendBtn *rod.Element
	var err error

	for _, selector := range sendSelectors {
		sendBtn, err = c.WaitForElementClickable(selector, 5)
		if err == nil && sendBtn != nil {
			c.LogInfo(fmt.Sprintf("找到发送按钮: %s", selector))
			break
		}
	}

	if sendBtn == nil {
		// 尝试查找发送图标
		sendBtn, err = c.Page.Element("button svg")
		if err != nil {
			return fmt.Errorf("未找到发送按钮")
		}
	}

	c.SleepMs(500)

	// 点击发送按钮
	if err := c.JSClick(sendBtn); err != nil {
		return fmt.Errorf("点击发送按钮失败: %v", err)
	}

	c.LogInfo("已点击发送按钮")
	c.SleepMs(2000)

	return nil
}

// waitForAnswer 等待并获取答案
func (c *DoubaoCollector) waitForAnswer() (string, error) {
	c.LogInfo("等待豆包回答...")

	timeout := 120 // 最大等待时间（秒）
	startTime := time.Now()
	lastAnswerLength := 0

	for time.Since(startTime).Seconds() < float64(timeout) {
		// 查找答案区域
		answerSelectors := []string{
			".message-content",
			".response-text",
			"[class*='assistant'] [class*='content']",
			"[class*='bot'] [class*='message']",
			".chat-message.bot",
			".answer-box",
		}

		for _, selector := range answerSelectors {
			answerElements, err := c.Page.Elements(selector)
			if err == nil && len(answerElements) > 0 {
				// 获取最后一个答案元素
				lastAnswer := answerElements[len(answerElements)-1]

				visible, _ := lastAnswer.Visible()
				if visible {
					text, err := lastAnswer.Text()
					if err == nil && len(strings.TrimSpace(text)) > 0 {
						// 检查是否正在生成
						isGenerating := strings.Contains(text, "正在") ||
							strings.Contains(text, "思考中") ||
							strings.Contains(text, "typing")

						if !isGenerating {
							// 检查答案是否还在增长
							currentLength := len(text)
							if currentLength == lastAnswerLength && currentLength > 10 {
								// 答案不再增长，认为已完成
								c.LogInfo("获取到完整答案")
								return strings.TrimSpace(text), nil
							}
							lastAnswerLength = currentLength
						}
					}
				}
			}
		}

		c.SleepMs(1500)
	}

	return "", fmt.Errorf("等待答案超时")
}

// SafeElement 安全地获取元素
func (c *DoubaoCollector) SafeElement(selector string) (*rod.Element, error) {
	exists, _, err := c.Page.Has(selector)
	if err != nil {
		return nil, err
	}
	if !exists {
		return nil, nil
	}
	return c.Page.Element(selector)
}