// Package sparkchat 是星火文本 websocket。会议总结不要用它。
//
// 海外 SDK 的文本模型是 domain=openai，见 sys/iflytek/openaichat。
// 这个包留着只因为星火协议和 openai 那条不是同一个地址。
package sparkchat

import (
	"bufio"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/url"
	"strings"
	"time"
)

const (
	defaultHost   = "spark-api.xf-yun.com"
	defaultDomain = "generalv3.5"
)

// Message 一轮对话。Role 取 system / user / assistant。
type Message struct {
	Role    string
	Content string
}

// Client 星火文本客户端。
type Client struct {
	AppID  string
	APIKey string
	Secret string
	Host   string
	Domain string
}

// New 用 SparkChain 凭据建客户端。
func New(appID, apiKey, secret, host, domain string) (*Client, error) {
	if strings.TrimSpace(appID) == "" || strings.TrimSpace(apiKey) == "" || strings.TrimSpace(secret) == "" {
		return nil, fmt.Errorf("iflytek spark: app_id/api_key/api_secret 都不能空")
	}
	host = strings.TrimSpace(host)
	if host == "" {
		host = defaultHost
	}
	domain = strings.TrimSpace(domain)
	if domain == "" {
		domain = defaultDomain
	}
	return &Client{AppID: appID, APIKey: apiKey, Secret: secret, Host: host, Domain: domain}, nil
}

// Chat 同步要一段回复。多帧内容拼在一起返回。
func (c *Client) Chat(messages []Message) (string, error) {
	if len(messages) == 0 {
		return "", fmt.Errorf("iflytek spark: 没有输入")
	}
	path := domainPath(c.Domain)
	conn, err := c.dial(path)
	if err != nil {
		return "", err
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(90 * time.Second))

	text := make([]map[string]string, 0, len(messages))
	for _, m := range messages {
		role := m.Role
		if role == "" {
			role = "user"
		}
		text = append(text, map[string]string{"role": role, "content": m.Content})
	}
	payload := map[string]any{
		"header": map[string]string{"app_id": c.AppID, "uid": "echomeet"},
		"parameter": map[string]any{
			"chat": map[string]any{
				"domain":      c.Domain,
				"temperature": 0.5,
				"max_tokens":  4096,
			},
		},
		"payload": map[string]any{
			"message": map[string]any{"text": text},
		},
	}
	raw, err := json.Marshal(payload)
	if err != nil {
		return "", err
	}
	if err := writeText(conn, raw); err != nil {
		return "", err
	}
	var b strings.Builder
	for {
		frame, err := readFrame(conn)
		if err != nil {
			if b.Len() > 0 && err == io.EOF {
				break
			}
			return "", err
		}
		if frame == nil {
			continue
		}
		var resp struct {
			Header struct {
				Code    int    `json:"code"`
				Message string `json:"message"`
				Status  int    `json:"status"`
			} `json:"header"`
			Payload struct {
				Choices struct {
					Status int `json:"status"`
					Text   []struct {
						Content string `json:"content"`
					} `json:"text"`
				} `json:"choices"`
			} `json:"payload"`
		}
		if json.Unmarshal(frame, &resp) != nil {
			continue
		}
		if resp.Header.Code != 0 {
			return "", fmt.Errorf("iflytek spark: code=%d %s", resp.Header.Code, resp.Header.Message)
		}
		for _, t := range resp.Payload.Choices.Text {
			b.WriteString(t.Content)
		}
		if resp.Header.Status == 2 || resp.Payload.Choices.Status == 2 {
			break
		}
	}
	out := strings.TrimSpace(b.String())
	if out == "" {
		return "", fmt.Errorf("iflytek spark: 空回复")
	}
	return out, nil
}

func domainPath(domain string) string {
	switch domain {
	case "general":
		return "/v1.1/chat"
	case "generalv2":
		return "/v2.1/chat"
	case "generalv3":
		return "/v3.1/chat"
	case "4.0Ultra":
		return "/v4.0/chat"
	default:
		return "/v3.5/chat"
	}
}

func (c *Client) dial(path string) (net.Conn, error) {
	date := time.Now().UTC().Format("Mon, 02 Jan 2006 15:04:05 GMT")
	origin := "host: " + c.Host + "\n" + "date: " + date + "\n" + "GET " + path + " HTTP/1.1"
	mac := hmac.New(sha256.New, []byte(c.Secret))
	_, _ = mac.Write([]byte(origin))
	signature := base64.StdEncoding.EncodeToString(mac.Sum(nil))
	authOrigin := fmt.Sprintf(`api_key="%s", algorithm="hmac-sha256", headers="host date request-line", signature="%s"`,
		c.APIKey, signature)
	auth := base64.StdEncoding.EncodeToString([]byte(authOrigin))
	q := url.Values{}
	q.Set("authorization", auth)
	q.Set("date", date)
	q.Set("host", c.Host)
	conn, err := tls.Dial("tcp", c.Host+":443", &tls.Config{ServerName: c.Host})
	if err != nil {
		return nil, fmt.Errorf("iflytek spark dial: %w", err)
	}
	key := make([]byte, 16)
	_, _ = rand.Read(key)
	secKey := base64.StdEncoding.EncodeToString(key)
	req := "GET " + path + "?" + q.Encode() + " HTTP/1.1\r\n" +
		"Host: " + c.Host + "\r\n" +
		"Upgrade: websocket\r\n" +
		"Connection: Upgrade\r\n" +
		"Sec-WebSocket-Key: " + secKey + "\r\n" +
		"Sec-WebSocket-Version: 13\r\n\r\n"
	if _, err := io.WriteString(conn, req); err != nil {
		conn.Close()
		return nil, err
	}
	br := bufio.NewReader(conn)
	status, err := br.ReadString('\n')
	if err != nil {
		conn.Close()
		return nil, err
	}
	if !strings.Contains(status, "101") {
		conn.Close()
		return nil, fmt.Errorf("iflytek spark handshake: %s", strings.TrimSpace(status))
	}
	for {
		line, err := br.ReadString('\n')
		if err != nil {
			conn.Close()
			return nil, err
		}
		if line == "\r\n" || line == "\n" {
			break
		}
	}
	return &bufferedConn{Conn: conn, r: br}, nil
}

// bufferedConn 把握手时多读进 bufio 的字节交还给后续读帧。
type bufferedConn struct {
	net.Conn
	r *bufio.Reader
}

func (c *bufferedConn) Read(p []byte) (int, error) { return c.r.Read(p) }

func writeText(conn net.Conn, payload []byte) error {
	mask := make([]byte, 4)
	_, _ = rand.Read(mask)
	var hdr []byte
	n := len(payload)
	switch {
	case n < 126:
		hdr = []byte{0x81, byte(n) | 0x80}
	case n < 65536:
		hdr = []byte{0x81, 126 | 0x80, byte(n >> 8), byte(n)}
	default:
		return fmt.Errorf("iflytek spark: 请求过长")
	}
	frame := append(hdr, mask...)
	masked := make([]byte, n)
	for i := 0; i < n; i++ {
		masked[i] = payload[i] ^ mask[i%4]
	}
	frame = append(frame, masked...)
	_, err := conn.Write(frame)
	return err
}

func readFrame(conn net.Conn) ([]byte, error) {
	for {
		hdr := make([]byte, 2)
		if _, err := io.ReadFull(conn, hdr); err != nil {
			return nil, err
		}
		opcode := hdr[0] & 0x0f
		masked := hdr[1]&0x80 != 0
		n := int(hdr[1] & 0x7f)
		switch n {
		case 126:
			ext := make([]byte, 2)
			if _, err := io.ReadFull(conn, ext); err != nil {
				return nil, err
			}
			n = int(ext[0])<<8 | int(ext[1])
		case 127:
			return nil, fmt.Errorf("iflytek spark: 帧过长")
		}
		var mask []byte
		if masked {
			mask = make([]byte, 4)
			if _, err := io.ReadFull(conn, mask); err != nil {
				return nil, err
			}
		}
		payload := make([]byte, n)
		if _, err := io.ReadFull(conn, payload); err != nil {
			return nil, err
		}
		if masked {
			for i := 0; i < n; i++ {
				payload[i] ^= mask[i%4]
			}
		}
		switch opcode {
		case 0x1, 0x0:
			return payload, nil
		case 0x8:
			return nil, io.EOF
		case 0x9: // ping
			pong := []byte{0x8A, 0}
			_, _ = conn.Write(pong)
			continue
		default:
			continue
		}
	}
}