package main

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

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)

		//----------------------------------
		//select kimi or deepseek
		kimi := 1
		if strings.HasPrefix(question, "kimi,") {
			kimi = 1
			question = strings.TrimPrefix(question, "kimi,")
		} else if strings.HasPrefix(question, "deepseek,") {
			kimi = 0
			question = strings.TrimPrefix(question, "deepseek,")
		}

		//----------------------------------
		//Call kimi or deepseek api
		myapikey := ""
		url := ""
		mymodel := ""
		if kimi == 1 {
			fmt.Println("Use kimi model")
			myapikey = "sk-gHXTd1mme286iRXrSX2cMedSzehS2lcx5I68vHXQcst8AxFh"
			url = "https://api.moonshot.cn/v1/chat/completions"
			//mymodel = "kimi-k2-0711-preview"
			mymodel = "moonshot-v1-8k"
		} else {
			fmt.Println("Use deekseek model")
			myapikey = "sk-6f2840ec32244476b0c661e952091db3"
			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 print_time() string {
	now := time.Now()
	return fmt.Sprintf("[%d-%02d-%02d %02d:%02d:%02d] ",
			now.Year(),
			now.Month(),
			now.Day(),
			now.Hour(),
			now.Minute(),
			now.Second())
}

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