package main

import (
	"net"
	"bufio"
	"fmt"
	"net/http"
	"encoding/json"
	"strings"
	"io"
)

type ChatResponse struct {
	Choices []struct {
		Delta struct {
			Content string `json:"content"`
		} `json:"delta"`
	} `json:"choices"`
}

func process(conn net.Conn) {
	defer conn.Close() // 关闭连接
	remoteAddr := conn.RemoteAddr().String()
	
	for {
		//----------------------------------
		//receive question from client
		reader := bufio.NewReader(conn)
		var buf [128]byte
		n, err := reader.Read(buf[:]) // 读取数据
		if err != nil {
			if err == io.EOF {
				//正常退出
				fmt.Println("client " + remoteAddr + " leaved! bye!")
			}else{
				//异常退出
				fmt.Println("read from client failed, err:", err)
			}
			break
		}
		recvStr := string(buf[:n])
		question := recvStr
		fmt.Println("问题：", question)

		//----------------------------------
		//call deepseek api
		myapikey := "<your-deepseek-api-key>"
		url := "https://api.deepseek.com/chat/completions"
		mymodel := "deepseek-chat"
		systemrole := "<系统角色描述>"
		
		payload := fmt.Sprintf(`{
                    "model": "%s",
                    "messages": [
                        {
                            "role": "system",
                            "content": "%s"
                        },
                        {
                            "role": "user",
                            "content": "%s"
                        }
                    ],
                    "stream": true,
                    "temperature": 0.3
                }`, mymodel, systemrole, question)

		//fmt.Print(payload)
		
		req, _ := http.NewRequest("POST", url, strings.NewReader(payload))
		req.Header.Add("Content-Type", "application/json")
		req.Header.Add("Authorization", "Bearer " + myapikey)

		res, _ := http.DefaultClient.Do(req)
		defer res.Body.Close()

		scanner := bufio.NewReader(res.Body)
		for {
			line, err := scanner.ReadString('\n')
			if err != nil {
				break
			}

			line = strings.TrimSpace(line)
			if line == "" || line == "data: [DONE]" {
				continue
			}
			if strings.HasPrefix(line, "data: ") {
				line = strings.TrimPrefix(line, "data: ")
			}
			
			var response ChatResponse
			if err := json.Unmarshal([]byte(line), &response); err != nil {
				continue
			}
			if len(response.Choices) > 0 {
				content := response.Choices[0].Delta.Content
				if content != "" {
					conn.Write([]byte(content))
				}
			}
		}

		//end of answer
		conn.Write([]byte("ENDENDEND\n"))
	}
}

func main() {
	listen, err := net.Listen("tcp", "0.0.0.0:8001")
	if err != nil {
		fmt.Println("listen failed, err:", err)
		return
	}
	fmt.Println("AI Server is up!")
	
	for {
		conn, err := listen.Accept() // 建立连接
		if err != nil {
			fmt.Println("accept failed, err:", err)
			continue
		}
		
		remoteAddr := conn.RemoteAddr().String()
		fmt.Println("Client " + remoteAddr + " is connected!")
		
		go process(conn) // 启动一个goroutine处理连接
	}
}