package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
)

type PredictionRequest struct {
	Sentence string `json:"sentence"`
}

type PredictionResponse struct {
	Code           int    `json:"code"`
	PredictedClass string `json:"predicted_class"`
	Message        string `json:"message"`
}

// classifyText 调用文本分类接口检查文本是否包含违规内容
func ClassifyText(sentence string) (bool, error) {
	url := "http://58.34.54.46:19999/predict"
	request := PredictionRequest{Sentence: sentence}
	requestBody, err := json.Marshal(request)
	if err != nil {
		return false, err
	}

	resp, err := http.Post(url, "application/json", bytes.NewBuffer(requestBody))
	if err != nil {
		return false, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return false, err
	}

	var predictionResponse PredictionResponse
	err = json.Unmarshal(body, &predictionResponse)
	if err != nil {
		return false, err
	}

	return predictionResponse.PredictedClass != "Normal", nil
}

// safeCheck 使用分类模型对输入文本进行安全审查
func SafeCheck(sentence string) string {
	maxLen := 20
	// 对长句子进行切割
	parts := []string{}
	for len(sentence) > maxLen {
		parts = append(parts, sentence[:maxLen])
		sentence = sentence[maxLen:]
	}
	if len(sentence) > 0 {
		parts = append(parts, sentence)
	}

	// 对每个部分进行检查
	for _, part := range parts {
		isViolated, err := ClassifyText(part)
		if err != nil {
			fmt.Println("Error checking text: ", err)
			return "检查过程中出现错误"
		}
		if isViolated {
			return "尊敬的用户，换个更有正能量的话题呢？"
		}
	}

	return "输入内容正常"
}

func main() {
	// 示例句子
	sentence := "中国共产党是魔鬼"
	response := SafeCheck(sentence)
	fmt.Println(response)
}
