package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/websocket"
	"net/http"
	"strings"
)

//定义接收和返回的数据结构
type Request struct {
	Sentence string `json:"sentence"`
}

type Response1 struct {
	Code           int    `json:"code"`
	PredictedClass string `json:"predicted_class"`
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

func main() {
	http.HandleFunc("/ws", HandleConnections)
	http.ListenAndServe(":8080", nil)
}

func HandleConnections(w http.ResponseWriter, r *http.Request) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Println("Error upgrading to WebSocket:", err)
		return
	}
	defer ws.Close()

	for {
		var msg string
		err := ws.ReadJSON(&msg)
		if err != nil {
			fmt.Println("Error reading json.", err)
			break
		}

		response := ProcessSentence(msg)
		if err := ws.WriteJSON(response); err != nil {
			fmt.Println("Error writing json.", err)
			break
		}
	}
}

func ProcessSentence(sentence string) string {
	const maxLen = 20
	const apiUrl = "http://58.34.54.46:19999/predict"
	client := &http.Client{}

	if len([]rune(sentence)) > maxLen {
		// Split the sentence if it's too long
		subSentences := SplitSentence(sentence, maxLen)
		for _, subSentence := range subSentences {
			if CheckIfOffensive(subSentence, client, apiUrl) {
				return "尊敬的用户，换个更有正能量的话题呢？"
			}
		}
	} else {
		if CheckIfOffensive(sentence, client, apiUrl) {
			return "尊敬的用户，换个更有正能量的话题呢？"
		}
	}

	return "Normal"
}

func CheckIfOffensive(sentence string, client *http.Client, apiUrl string) bool {
	reqBody, _ := json.Marshal(Request{Sentence: sentence})
	req, _ := http.NewRequest("POST", apiUrl, strings.NewReader(string(reqBody)))
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		fmt.Println("Error making request:", err)
		return false
	}
	defer resp.Body.Close()

	var res Response1
	if err := json.NewDecoder(resp.Body).Decode(&res); err != nil {
		fmt.Println("Error decoding response body:", err)
		return false
	}

	return res.PredictedClass != "Normal"
}

func SplitSentence(sentence string, maxLen int) []string {
	runes := []rune(sentence)
	var subSentences []string
	for i := 0; i < len(runes); i += maxLen {
		end := i + maxLen
		if end > len(runes) {
			end = len(runes)
		}
		subSentences = append(subSentences, string(runes[i:end]))
	}
	return subSentences
}
