package ai

import (
	"crypto/hmac"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/zeromicro/go-zero/core/logx"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

/**
 *  WebAPI 接口调用示例 接口文档（必看）：https://www.xfyun.cn/doc/spark/Web.html
 * 错误码链接：https://www.xfyun.cn/doc/spark/%E6%8E%A5%E5%8F%A3%E8%AF%B4%E6%98%8E.html（code返回错误码时必看）
 * @author iflytek
 */

var (
	hostChatUrl         = "wss://spark-api.xf-yun.com/v3.5/chat"
	hostFileUploadUrl   = "https://chatdoc.xfyun.cn/openapi/v1/file/upload"
	hostChatDocumentUrl = "wss://chatdoc.xfyun.cn/openapi/chat"
)

func ChatSpark(appid string, apiKey string, apiSecret string, question string, channel chan string) (answer string, err error) {
	// fmt.Println(HmacWithShaTobase64("hmac-sha256", "hello\nhello", "hello"))
	// st := time.Now()
	d := websocket.Dialer{
		HandshakeTimeout: 5 * time.Second,
	}
	conn, resp, err := d.Dial(assembleAuthUrl(hostChatUrl, apiKey, apiSecret), 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 := genParams1(appid, question)
		conn.WriteJSON(data)

	}()

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

		var data SparkChatMessage
		err = json.Unmarshal(msg, &data)
		if err != nil {
			logx.Error("Error parsing JSON:", err)
			return
		}
		logx.Info(string(msg))
		if data.Header.Code != 0 {
			logx.Error("ws error code ", data.Header.Code)
			return
		}
		status := data.Payload.Choices.Status
		text := data.Payload.Choices.Text
		var content string
		if len(text) > 0 {
			content = text[0].Content
			channel <- content
		}
		if status != 2 {
			answer += content
		} else {
			answer += content
			totalTokens := data.Payload.Usage.Text.TotalTokens
			logx.Infof("收到最终结果 total_tokens: %v answer:%v", totalTokens, answer)
			conn.Close()
			break
		}
	}
	return answer, nil
}

type SparkChatMessage struct {
	Header struct {
		Code    int    `json:"code"`
		Message string `json:"message"`
		Sid     string `json:"sid"`
		Status  int    `json:"status"`
	} `json:"header"`
	Payload struct {
		Choices struct {
			Status int `json:"status"`
			Seq    int `json:"seq"`
			Text   []struct {
				Content string `json:"content"`
				Role    string `json:"role"`
				Index   int    `json:"index"`
			} `json:"text"`
		} `json:"choices"`
		Usage struct {
			Text struct {
				QuestionTokens   int `json:"question_tokens"`
				PromptTokens     int `json:"prompt_tokens"`
				CompletionTokens int `json:"completion_tokens"`
				TotalTokens      int `json:"total_tokens"`
			} `json:"text"`
		} `json:"usage"`
	} `json:"payload"`
}

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

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

	data := map[string]interface{}{ // 根据实际情况修改返回的数据结构和字段名
		"header": map[string]interface{}{ // 根据实际情况修改返回的数据结构和字段名
			"app_id": appid, // 根据实际情况修改返回的数据结构和字段名
		},
		"parameter": map[string]interface{}{ // 根据实际情况修改返回的数据结构和字段名
			"chat": map[string]interface{}{ // 根据实际情况修改返回的数据结构和字段名
				"domain":      "generalv3.5", // 根据实际情况修改返回的数据结构和字段名
				"temperature": float64(0.8),  // 根据实际情况修改返回的数据结构和字段名
				"top_k":       int64(6),      // 根据实际情况修改返回的数据结构和字段名
				"max_tokens":  int64(2048),   // 根据实际情况修改返回的数据结构和字段名
				"auditing":    "default",     // 根据实际情况修改返回的数据结构和字段名
			},
		},
		"payload": map[string]interface{}{ // 根据实际情况修改返回的数据结构和字段名
			"message": map[string]interface{}{ // 根据实际情况修改返回的数据结构和字段名
				"text": messages, // 根据实际情况修改返回的数据结构和字段名
			},
		},
	}
	return data // 根据实际情况修改返回的数据结构和字段名
}

// 创建鉴权url  apikey 即 hmac username
func assembleAuthUrl(hosturl string, appKey, apiSecret string) string {
	ul, err := url.Parse(hosturl)
	if err != nil {
		fmt.Println(err)
	}
	//签名时间
	date := time.Now().UTC().Format(time.RFC1123)
	//date = "Tue, 28 May 2019 09:10:42 MST"
	//参与签名的字段 host ,date, request-line
	signString := []string{"host: " + ul.Host, "date: " + date, "GET " + ul.Path + " HTTP/1.1"}
	//拼接签名字符串
	sgin := strings.Join(signString, "\n")
	// fmt.Println(sgin)
	//签名结果
	sha := HmacWithSha256base64(sgin, apiSecret)
	// fmt.Println(sha)
	//构建请求参数 此时不需要urlencoding
	authUrl := fmt.Sprintf(`api_key="%v", algorithm="hmac-sha256", headers="host date request-line", signature="%v"`, appKey, sha)
	//将请求参数使用base64编码
	authorization := base64.StdEncoding.EncodeToString([]byte(authUrl))

	v := url.Values{}
	v.Add("host", ul.Host)
	v.Add("date", date)
	v.Add("authorization", authorization)
	//将编码后的字符串url encode后添加到url后面
	callurl := hosturl + "?" + v.Encode()
	return callurl
}

func HmacWithSha256base64(data, key string) string {
	mac := hmac.New(sha256.New, []byte(key))
	mac.Write([]byte(data))
	encodeData := mac.Sum(nil)
	return base64.StdEncoding.EncodeToString(encodeData)
}

func HmacWithSha1base64(data, key string) string {
	mac := hmac.New(sha1.New, []byte(key))
	mac.Write([]byte(data))
	encodeData := mac.Sum(nil)
	return base64.StdEncoding.EncodeToString(encodeData)
}

func readResp(resp *http.Response) string {
	if resp == nil {
		return ""
	}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("code=%d,body=%s", resp.StatusCode, string(b))
}

type SparkUserMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}