package tests

import (
	"apigo.cc/ai/ai"
	"apigo.cc/ai/ai/js"
	"apigo.cc/ai/ai/llm"
	"encoding/hex"
	"fmt"
	"github.com/ssgo/u"
	"testing"
)

func TestAgent(t *testing.T) {
	ai.Init()
	//testChat(t, "gpt")
	// 尚未支持
	//testCode(t, "gpt")
	// keya 消耗过大，计费不合理
	//testAskWithImage(t, "gpt", "4032.jpg")
	// keya 不支持
	//testMakeImage(t, "gpt", "冬天大雪纷飞，一个男人身穿军绿色棉大衣，戴着红色围巾和绿色帽子走在铺面大雪的小镇路上", "")

	//testChat(t, "glm")
	//testCode(t, "glm", "", keys.Zhipu)
	//testSearch(t, "glm", "", keys.Zhipu)

	// 测试图片识别
	//testAskWithImage(t, "glm", "", keys.Zhipu, "4032.jpg")

	// 视频似乎尚不支持 glm-4v-plus
	//testAskWithVideo(t, "glm", "", keys.Zhipu, "glm-4v", "1080.mp4")

	//testMakeImage(t, "glm", "", keys.Zhipu, "冬天大雪纷飞，一个男人身穿军绿色棉大衣，戴着红色围巾和绿色帽子走在铺面大雪的小镇路上", "")
	//testMakeVideo(t, "glm", "", keys.Zhipu, "大雪纷飞，男人蹦蹦跳跳", "https://aigc-files.bigmodel.cn/api/cogview/20240904133130c4b7121019724aa3_0.png")

	//testJS(t)
	//testFile(t)
	testUtil(t)
}

func testChat(t *testing.T, llmName string) {
	lm := llm.Get(llmName)

	if lm == nil {
		t.Fatal("agent is nil")
	}

	r, usage, err := lm.FastAsk(llm.Messages().User().Text("你是什么模型，请给出具体名称、版本号").Make(), func(text string) {
		fmt.Print(u.BCyan(text))
		fmt.Print(" ")
	})

	if err != nil {
		t.Fatal(err)
	}

	fmt.Println()
	fmt.Println("result:", r)
	fmt.Println("usage:", u.JsonP(usage))
}

func testCode(t *testing.T, llmName string) {
	lm := llm.Get(llmName)

	if lm == nil {
		t.Fatal("agent is nil")
	}

	r, usage, err := lm.CodeInterpreterAsk(llm.Messages().User().Text("计算[5,10,20,700,99,310,978,100]的平均值和方差。").Make(), func(text string) {
		fmt.Print(u.BCyan(text))
		fmt.Print(" ")
	})

	if err != nil {
		t.Fatal(err)
	}

	fmt.Println()
	fmt.Println("result:", r)
	fmt.Println("usage:", u.JsonP(usage))
}

func testSearch(t *testing.T, llmName string) {
	lm := llm.Get(llmName)

	if lm == nil {
		t.Fatal("agent is nil")
	}

	r, usage, err := lm.WebSearchAsk(llm.Messages().User().Text("今天上海的天气怎么样？").Make(), func(text string) {
		fmt.Print(u.BCyan(text))
		fmt.Print(" ")
	})

	if err != nil {
		t.Fatal(err)
	}

	fmt.Println()
	fmt.Println("result:", r)
	fmt.Println("usage:", u.JsonP(usage))
}

func testAskWithImage(t *testing.T, llmName, imageFile string) {
	lm := llm.Get(llmName)

	if lm == nil {
		t.Fatal("agent is nil")
	}

	ask := `请回答：
1、我用的是什么操作系统？
2、现在的日期时间是？请给出年月日。
3、正在用什么软件播放什么歌？谁演唱的？歌曲的大意是？
4、后面的浏览器中正在浏览什么内容？猜测一下我浏览这个网页是想干嘛？
`
	r, usage, err := lm.MultiAsk(llm.Messages().User().Text(ask).Image("data:image/jpeg;base64,"+u.Base64(u.ReadFileBytesN(imageFile))).Make(), func(text string) {
		fmt.Print(u.BCyan(text))
		fmt.Print(" ")
	})

	if err != nil {
		t.Fatal("发生错误", err.Error())
	}

	fmt.Println()
	fmt.Println("result:", r)
	fmt.Println("usage:", u.JsonP(usage))
}

func testAskWithVideo(t *testing.T, llmName, videoFile string) {
	lm := llm.Get(llmName)

	if lm == nil {
		t.Fatal("agent is nil")
	}

	ask := `请回答：
1、正在用什么软件播放什么歌？谁演唱的？歌曲的大意是？
4、后面的浏览器中正在浏览什么内容？猜测一下我浏览这个网页是想干嘛？
`

	r, usage, err := lm.MultiAsk(llm.Messages().User().Text(ask).Video("data:video/mp4,"+u.Base64(u.ReadFileBytesN(videoFile))).Make(), func(text string) {
		fmt.Print(u.BCyan(text))
		fmt.Print(" ")
	})

	if err != nil {
		t.Fatal("发生错误", err.Error())
	}

	fmt.Println()
	fmt.Println("result:", r)
	fmt.Println("usage:", u.JsonP(usage))
}

func testMakeImage(t *testing.T, llmName, prompt, refImage string) {
	lm := llm.Get(llmName)

	if lm == nil {
		t.Fatal("agent is nil")
	}

	r, err := lm.FastMakeImage(prompt, llm.GCConfig{
		Size: "1024x1024",
	})

	if err != nil {
		t.Fatal("发生错误", err.Error())
	}

	fmt.Println()
	for i, v := range r {
		fmt.Println("result:", i, v)
	}
}

func testMakeVideo(t *testing.T, llmName, prompt, refImage string) {
	lm := llm.Get(llmName)

	if lm == nil {
		t.Fatal("agent is nil")
	}

	r, covers, err := lm.FastMakeVideo(prompt, llm.GCConfig{
		Size: "1280x720",
		Ref:  refImage,
	})

	if err != nil {
		t.Fatal("发生错误", err.Error())
	}

	fmt.Println()
	for i, v := range r {
		fmt.Println("result:", i, v, covers[i])
	}
}

func testJS(t *testing.T) {
	r1, err := ai.RunFile("test.js", "1+2=4吗")
	if err != nil {
		t.Fatal("发生错误", err.Error())
	}
	r := js.ChatResult{}
	u.Convert(r1, &r)

	fmt.Println()
	fmt.Println("result:", r.Result)
	fmt.Println("usage:", u.JsonP(r.TokenUsage))
}

func testFile(t *testing.T) {
	r1, err := ai.Run(`
import fs from './lib/file'
import out from './lib/console'
let r = fs.read('test.js')
return r
`, "test.js")
	if err != nil {
		t.Fatal("发生错误", err.Error())
	}

	fmt.Println()
	fmt.Println("result:", r1)
}

func testUtil(t *testing.T) {
	fmt.Println(hex.EncodeToString(u.Sha256([]byte{0, 98, 2})))
	r, err := ai.Run(`
let s = String.fromCharCode(0, 98, 2)
//let r = util.base64(s)
//let r2 = util.unBase64(r)
//console.info(r)
//console.info(r2.length, r2)
//console.info(s.charCodeAt(0), r2.charCodeAt(0))
//return r2 === s
let r = util.hex(util.sha256(s))
console.info(r)
return r == '278f25d5e7c09b2cdde1ed72dd83fbb7a7d7dd668ef967c9af65f68aa04049cd'
`, "test.js")
	if err != nil {
		t.Fatal("发生错误", err.Error())
	}
	if r != true {
		t.Fatal("运行结果不正确", r)
	}

	fmt.Println()
	fmt.Println("result:", r)
}