package ai

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/websocket"
	"github.com/zeromicro/go-zero/core/jsonx"
	"github.com/zeromicro/go-zero/core/logx"
	"time"
)

func ChatSparkDocument(appid string, apiKey string, apiSecret string, fileIds []string, question string, channel chan string) (answer string, err error) {
	d := websocket.Dialer{
		HandshakeTimeout: 5 * time.Second,
	}
	//握手并建立websocket 连接
	ts := fmt.Sprintf("%d", time.Now().Unix())
	conn, resp, err := d.Dial(hostChatDocumentUrl+"?"+"appId="+appid+"&timestamp="+ts+"&signature="+getSignature(appid, apiSecret, ts), nil)
	if err != nil {
		logx.Error(readResp(resp) + err.Error())
		return
	} else if resp.StatusCode != 101 {
		logx.Error(readResp(resp) + err.Error())
		return
	}

	go func() {
		data := genDocumentChatParams(fileIds, question)
		logx.Info(jsonx.MarshalToString(data))
		conn.WriteJSON(data)

	}()

	//获取返回的数据
	for {
		var msg []byte
		_, msg, err = conn.ReadMessage()
		if err != nil {
			logx.Error("read message error:", err)
			break
		}

		var data SparkDocumentChatResponseMessage
		err = json.Unmarshal(msg, &data)
		if err != nil {
			logx.Error("Error parsing JSON:", err)
			return
		}
		logx.Info(string(msg))
		// 异常时，错误信息作为回答
		if data.Code != 0 {
			channel <- data.Content
			answer = data.Content
			return
		}
		content := data.Content
		channel <- content
		if data.Status != 2 {
			answer += content
		} else {
			answer += content
			logx.Infof("收到最终结果 answer:%v", answer)
			conn.Close()
			break
		}
	}
	return answer, nil
}

/*
若返回 fileRefer 为空，提示 "抱款，在文档中没有找到与提问相关的内容，请尝试换个问题问问吧。"表示提问未匹配到文档内容，可以降低chatExtends.wikiFilterScore以降低匹配阈值，也可以开启chatExtends.sparkWhenWithoutEmbedding用大模型兜底
*/
type SparkDocumentChatResponseMessage struct {
	Code      int    `json:"code"`      // 错误码 ，0 标识成功
	Content   string `json:"content"`   // 回答内容
	FileRefer string `json:"fileRefer"` // 文档引用，status=99 的时候返回；结构是个 Map，key=文件 id，value=引用的文段列表（对应 fileTrunks 的 index
	Status    int    `json:"status"`    // 会话状态，取值为[0,1,2,99]；0 代表首次结果；1 代表中间结果；2 代表最后一个结果；99 代表引用的文档及文段
}

// 生成参数
func genDocumentChatParams(fileIds []string, question string) map[string]interface{} { // 根据实际情况修改返回的数据结构和字段名

	messages := []Message{
		{Role: "user", Content: question},
	}

	data := map[string]interface{}{ // 根据实际情况修改返回的数据结构和字段名
		"chatExtends": map[string]interface{}{ // 根据实际情况修改返回的数据结构和字段名
			"wikiPromptTpl":   "请将以下内容作为已知信息：\n<wikicontent>\n请根据以上内容回答用户的问题。\n问题:<wikiquestion>\n回答:", // 根据实际情况修改返回的数据结构和字段名
			"wikiFilterScore": 0.82,
			"temperature":     0.5,
		},
		"fileIds":  fileIds,
		"messages": messages,
	}
	return data // 根据实际情况修改返回的数据结构和字段名
}