geoGo/internal/collect/wenxin.go

package collect

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

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

// WenxinCollector 文心一言收集器
type WenxinCollector struct {
	*BaseCollector
}

// NewWenxinCollector 创建文心一言收集器
func NewWenxinCollector(ctx context.Context, params *CollectParams, cfg *config.Config, logger *log.Logger) CollectorInterface {
	collector := &WenxinCollector{
		BaseCollector: NewBaseCollector(ctx, params, cfg, logger),
	}

	// 设置文心一言的URL
	collector.LoginURL = "https://passport.baidu.com/v2/?login"
	collector.ChatURL = "https://yiyan.baidu.com/"

	return collector
}

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

	// 如果在登录页面，说明未登录
	if strings.Contains(currentURL, "passport.baidu.com") {
		return false
	}

	// 检查页面上是否存在内容为"登录"或"Login"的button，如果存在说明未登录
	loginButtons, err := c.Page.Elements("button")
	if err == nil {
		for _, btn := range loginButtons {
			text, _ := btn.Text()
			trimmedText := strings.TrimSpace(text)
			if trimmedText == "登录" || trimmedText == "Login" {
				c.LogInfo(fmt.Sprintf("检测到页面上有'%s'按钮，说明未登录", trimmedText))
				return false
			}
		}
	}

	// 如果没有找到"登录"或"Login"按钮，说明已登录
	return true
}

// WaitLogin 等待登录
func (c *WenxinCollector) WaitLogin() (bool, string) {
	c.LogInfo("开始等待文心一言登录...")

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

	// 访问聊天页面
	c.Page.MustNavigate(c.ChatURL)
	c.Sleep(3)

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

	c.LogInfo("检测到未登录，请在当前页面完成登录（扫码或输入账号密码）...")

	// 不跳转页面，在当前页面循环检查登录按钮是否存在
	// 最多等待300秒
	for i := 0; i < 5000; i++ {
		// 检查页面上是否还存在"登录"或"Login"按钮
		loginButtonExists := false
		buttons, err := c.Page.Elements("button")
		if err == nil {
			for _, btn := range buttons {
				text, _ := btn.Text()
				trimmedText := strings.TrimSpace(text)
				if trimmedText == "登录" || trimmedText == "Login" {
					loginButtonExists = true
					break
				}
			}
		}

		// 如果登录按钮不存在，说明已登录
		if !loginButtonExists {
			c.Sleep(2) // 等待页面稳定
			c.SaveCookies()
			c.LogInfo("登录成功：登录按钮已消失")
			return true, "login_success"
		}

		// 每秒检查一次
		time.Sleep(1 * time.Second)

		// 每30秒输出一次提示
		if i > 0 && i%30 == 0 {
			c.LogInfo(fmt.Sprintf("等待登录中... 已等待 %d 秒", i))
		}
	}

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

// AskQuestion 提问并获取答案
func (c *WenxinCollector) 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()

	//初始化页面（加载cookies和检查登录）
	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 *WenxinCollector) inputQuestion(question string) error {
	c.LogInfo("输入问题...")

	// 文心一言的输入框选择器 - 根据实际页面结构调整
	inputSelectors := []string{
		"[contenteditable='true']",
		"div[contenteditable]",
		".editable__T7WAW4uW",
		"[class*='editable']",
	}

	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)

	// fallback: 使用Focus + Input
	inputBox.Focus()
	c.SleepMs(200)
	inputBox.Input(question)

	c.LogInfo(fmt.Sprintf("问题已输入: %s", question))
	c.SleepMs(1000)

	return nil
}

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

	// 使用正则匹配包含"send"的class（防CSS混淆）
	allElements, err := c.Page.Elements("*")
	if err != nil {
		return fmt.Errorf("获取页面元素失败: %v", err)
	}

	var sendBtn *rod.Element
	for _, elem := range allElements {
		classAttr, _ := elem.Attribute("class")
		if classAttr != nil && strings.Contains(strings.ToLower(*classAttr), "send") {
			// 检查是否是可点击的元素（button、div等）
			tagName, _ := elem.Property("tagName")
			if tagName.Str() == "BUTTON" || tagName.Str() == "DIV" {
				sendBtn = elem
				c.LogInfo(fmt.Sprintf("通过正则找到发送按钮: class=%s, tag=%s", *classAttr, tagName.Str()))
				break
			}
		}
	}

	if sendBtn == nil {
		// fallback: 尝试查找最后一个button
		buttons, _ := c.Page.Elements("button")
		if len(buttons) > 0 {
			sendBtn = buttons[len(buttons)-1]
			c.LogInfo("使用最后一个button作为发送按钮")
		}
	}

	if sendBtn == nil {
		return fmt.Errorf("未找到发送按钮")
	}

	c.SleepMs(500)

	// 滚动到可见区域
	if err := sendBtn.ScrollIntoView(); err != nil {
		c.LogInfo(fmt.Sprintf("滚动失败: %v", err))
	}
	c.SleepMs(300)

	// 点击发送按钮
	c.LogInfo("执行点击...")
	if err := sendBtn.Click(proto.InputMouseButtonLeft, 1); err != nil {
		return fmt.Errorf("点击发送按钮失败: %v", err)
	}

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

	// 检测是否发送成功：检查send按钮是否消失或变成pause按钮
	maxWaitTime := 10 // 最多等待10秒
	for i := 0; i < maxWaitTime*2; i++ {
		// 检查是否存在pause开头的按钮（表示正在生成）
		pauseExists, err := c.hasPauseButton()
		if err == nil && pauseExists {
			c.LogInfo("✓ 检测到pause按钮，消息发送成功，AI正在回答...")
			return nil
		}

		// 检查send按钮是否还存在
		sendExists, _ := c.hasSendButton()
		if !sendExists {
			c.LogInfo("✓ send按钮已消失，消息发送成功")
			return nil
		}

		c.SleepMs(500)
	}

	c.LogInfo("⚠ 无法确认消息是否发送成功，但已尽力尝试")
	return nil
}

// hasSendButton 检查是否存在send开头的按钮
func (c *WenxinCollector) hasSendButton() (bool, error) {
	allElements, err := c.Page.Elements("*")
	if err != nil {
		return false, err
	}

	for _, elem := range allElements {
		classAttr, _ := elem.Attribute("class")
		if classAttr != nil && strings.Contains(strings.ToLower(*classAttr), "send") {
			tagName, _ := elem.Property("tagName")
			if tagName.Str() == "BUTTON" || tagName.Str() == "DIV" {
				return true, nil
			}
		}
	}
	return false, nil
}

// hasPauseButton 检查是否存在pause开头的按钮
func (c *WenxinCollector) hasPauseButton() (bool, error) {
	allElements, err := c.Page.Elements("*")
	if err != nil {
		return false, err
	}

	for _, elem := range allElements {
		classAttr, _ := elem.Attribute("class")
		if classAttr != nil && strings.Contains(strings.ToLower(*classAttr), "pause") {
			tagName, _ := elem.Property("tagName")
			if tagName.Str() == "BUTTON" || tagName.Str() == "DIV" {
				return true, nil
			}
		}
	}
	return false, nil
}

// waitForAnswer 等待并获取答案（处理流式输出）
func (c *WenxinCollector) waitForAnswer() (string, error) {
	c.LogInfo("等待AI回答...")

	timeout := 180 // 最大等待时间（秒），流式输出可能需要更长时间
	startTime := time.Now()

	var lastAnswer string
	var stableCount int           // 稳定计数器，连续N次内容不变则认为完成
	const requiredStableCount = 5 // 需要连续5次内容不变才认为完成
	isAnswering := false          // 标记是否正在回答中

	for time.Since(startTime).Seconds() < float64(timeout) {
		// 检查是否存在pause按钮（表示正在生成答案）
		pauseExists, _ := c.hasPauseButton()

		if pauseExists {
			if !isAnswering {
				c.LogInfo("检测到pause按钮，AI正在生成回答...")
				isAnswering = true
			}
		} else if isAnswering {
			// pause按钮消失，可能回答完成了
			c.LogInfo("pause按钮消失，检查回答是否完成...")

			// 再等待几次确认内容稳定
			if stableCount >= requiredStableCount && lastAnswer != "" {
				c.LogInfo(fmt.Sprintf("✓ AI回答完成，最终长度: %d 字符", len(lastAnswer)))
				return lastAnswer, nil
			}
		}

		// 直接通过ID查找答案容器
		answerElem, err := c.Page.Element("#answer_text_id")
		var answerText string

		if err == nil && answerElem != nil {
			// 获取整个HTML内容
			htmlContent, err := answerElem.HTML()
			if err == nil && len(strings.TrimSpace(htmlContent)) > 30 {
				// 清理HTML标签，只保留纯文本
				answerText = CleanHTMLTags(htmlContent)
				c.LogInfo(fmt.Sprintf("找到答案容器，清理后文本长度: %d", len(answerText)))
			} else {
				// 如果HTML获取失败，尝试获取文本
				textContent, _ := answerElem.Text()
				answerText = strings.TrimSpace(textContent)
				c.LogInfo(fmt.Sprintf("找到答案容器，文本长度: %d", len(answerText)))
			}
		} else {
			c.LogInfo("未找到#answer_text_id元素")
		}

		// 检查是否获取到答案
		if answerText != "" && len(answerText) > 30 {
			// 检查内容是否稳定（流式输出完成）
			if answerText == lastAnswer {
				stableCount++
				c.LogInfo(fmt.Sprintf("答案稳定中... (%d/%d), 长度: %d", stableCount, requiredStableCount, len(answerText)))

				// 如果pause按钮不存在且内容稳定，说明回答完成
				if !pauseExists && stableCount >= requiredStableCount {
					c.LogInfo(fmt.Sprintf("✓ AI回答完成，最终长度: %d 字符", len(answerText)))
					return answerText, nil
				}
			} else {
				// 内容还在变化，重置计数器
				stableCount = 0
				lastAnswer = answerText
				if pauseExists {
					c.LogInfo(fmt.Sprintf("检测到流式输出，当前长度: %d 字符", len(answerText)))
				}
			}
		}

		c.SleepMs(1500) // 每1.5秒检查一次

		// 每10秒输出一次等待状态
		elapsed := int(time.Since(startTime).Seconds())
		if elapsed > 0 && elapsed%10 == 0 {
			c.LogInfo(fmt.Sprintf("等待AI回答中... 已等待 %d 秒", elapsed))
		}
	}

	return "", fmt.Errorf("等待答案超时（%d秒）", timeout)
}

// SafeElement 安全地获取元素
func (c *WenxinCollector) 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)
}