package openai_test

import (
	"context"
	"encoding/json"
	"fmt"
	"testing"
	"yunyan/sys/openai"
	tavily "yunyan/sys/websearch/tavilysearch"

	ai "github.com/sashabaranov/go-openai"
)

type Tool_Search_Req struct {
	Query string `json:"query"`
}
type search struct{}

func (this *search) GetToolInfo() ai.Tool {
	// 定义工具列表
	tool := ai.Tool{
		Type: "function",
		Function: &ai.FunctionDefinition{
			Name:        "search",
			Description: "通过Tavily搜索引擎获取实时网络信息,例如查询股票id,实时新闻等",
			Parameters: openai.Parameters{
				Type: "object",
				Properties: map[string]openai.Property{
					"query": {
						Type:        "string",
						Description: "搜索关键词",
					},
				},
				Required: []string{"query"},
			},
		},
	}
	return tool
}
func (this *search) Execute(args json.RawMessage) (data any, text string, err error) {
	var (
		req    Tool_Search_Req
		result *tavily.TavilyResponse
	)
	fmt.Println("search", string(args))
	if err := json.Unmarshal(args, &req); err != nil {
		return nil, "输入参数错误", fmt.Errorf("参数解析失败: %v", err)
	}
	if result, err = tavily.Search(context.Background(), req.Query, 10, 24); err != nil {
		return
	}
	data = result
	text = result.Results[0].Content
	return
}

func Test_Sys_Chat(t *testing.T) {
	if sys, err := openai.NewSys(
		openai.SetBaseURL("https://api.deepseek.com/v1"),
		openai.SetToken("sk-3adfd188a3134e718bbf704f525aff17"),
		openai.SetModel("deepseek-chat"),
	); err != nil {
		fmt.Printf("Sys Init err:%v", err)
	} else {
		var resp *openai.ChatResponse
		// 初始化对话
		messages := []openai.ChatReq{
			{
				Role:    ai.ChatMessageRoleUser,
				Content: "你好呀",
			},
		}
		resp, err = sys.Chat(context.Background(), messages)
		fmt.Printf("resp:%+v err:%v", resp, err)
	}
}

func Test_Sys_ChatSteam(t *testing.T) {
	tavily.OnInit(nil,
		tavily.SetApiKey("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"),
	)
	if sys, err := openai.NewSys(
		openai.SetBaseURL("https://api.deepseek.com/v1"),
		openai.SetToken("sk-3adfd188a3134e718bbf704f525aff17"),
		openai.SetModel("deepseek-chat"),
		openai.SetMaxfunccall(5),
	); err != nil {
		fmt.Printf("Sys Init err:%v", err)
	} else {
		// 初始化对话
		messages := []openai.ChatReq{
			{
				Role:    ai.ChatMessageRoleUser,
				Content: "茅台的今日股价如何?",
			},
		}
		sys.RegisterTools(&search{})
		choiceChan := make(chan *openai.ChatResponseChoice, 1)
		go sys.ChatForSteams(context.Background(), messages, choiceChan)
		for v := range choiceChan {
			fmt.Printf("%+v\n", v)
		}
	}
}

func Test_Sys_doubao(t *testing.T) {
	if sys, err := openai.NewSys(
		openai.SetBaseURL("https://ark.cn-beijing.volces.com/api/v3/bots"),
		openai.SetToken("418ec475-e2dc-4b76-8aca-842d81bc3466"),
		openai.SetModel("bot-20250328144214-rrln6"),
	); err != nil {
		fmt.Printf("Sys Init err:%v", err)
	} else {
		var resp *openai.ChatResponse
		// 初始化对话
		messages := []openai.ChatReq{
			{
				Role:    ai.ChatMessageRoleUser,
				Content: "茅台今日股价？",
			},
		}
		resp, err = sys.Chat(context.Background(), messages)
		fmt.Printf("resp:%+v err:%v", resp, err)
	}
}
func Test_Sys_doubao_steam(t *testing.T) {
	tavily.OnInit(nil,
		tavily.SetApiKey("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"),
	)
	if sys, err := openai.NewSys(
		openai.SetBaseURL("https://ark.cn-beijing.volces.com/api/v3/bots"),
		openai.SetToken("418ec475-e2dc-4b76-8aca-842d81bc3466"),
		openai.SetModel("bot-20250328144214-rrln6"),
	); err != nil {
		fmt.Printf("Sys Init err:%v", err)
	} else {
		// 初始化对话
		messages := []openai.ChatReq{
			{
				Role:    ai.ChatMessageRoleUser,
				Content: "茅台的今日股价如何?",
			},
		}
		choiceChan := make(chan *openai.ChatResponseChoice, 1)
		go sys.ChatForSteams(context.Background(), messages, choiceChan)
		for v := range choiceChan {
			fmt.Printf("%+v\n", v)
		}
	}
}

func Test_Sys_oneapi_steam(t *testing.T) {
	// tavily.OnInit(nil,
	// 	tavily.SetApiKey("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"),
	// )
	if sys, err := openai.NewSys(
		openai.SetBaseURL("http://127.0.0.1:3000/v1"),
		openai.SetToken("sk-gMpCSTxYSpMoMHXz25B9Bc6598Bd431b8cCe6996625c8fD1"),
		openai.SetModel("gpt-3.5-turbo"),
	); err != nil {
		fmt.Printf("Sys Init err:%v", err)
	} else {
		// 初始化对话
		messages := []openai.ChatReq{
			{
				Role:    ai.ChatMessageRoleUser,
				Content: "股票代码:sh601009,今日行情咋样",
			},
		}
		choiceChan := make(chan *openai.ChatResponseChoice, 1)
		// sys.RegisterTools(&search{})
		go sys.ChatForSteams(context.Background(), messages, choiceChan)
		for v := range choiceChan {
			fmt.Printf("%+v\n", v)
		}
	}
}