package main

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

// 定義必要的數(shù)據(jù)結(jié)構(gòu)
type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type RequestBody struct {
	Model       string    `json:"model"`
	Messages    []Message `json:"messages"`
	Temperature float64   `json:"temperature"`
	Stream      bool      `json:"stream"`
}

type Choice struct {
	Delta struct {
		Content string `json:"content"`
	} `json:"delta"`
}

type ResponseBody struct {
	Choices []Choice `json:"choices"`
}

const (
	apiURL      = "https://api.example.com/v1/chat/completions" // 替換為實(shí)際的 API 地址
	authToken   = "your-auth-token"                             // 替換為實(shí)際的 Token
	model       = "gpt-3.5-turbo"
	temperature = 0.7
)

func StreamHandler(w http.ResponseWriter, r *http.Request) {
	// 從查詢參數(shù)獲取輸入內(nèi)容
	content := r.URL.Query().Get("content")
	if content == "" {
		http.Error(w, "Missing 'content' parameter", http.StatusBadRequest)
		return
	}

	// 構(gòu)造請(qǐng)求體
	message := Message{
		Role:    "user",
		Content: content,
	}
	requestBody := RequestBody{
		Model:       model,
		Messages:    []Message{message},
		Temperature: temperature,
		Stream:      true,
	}
	jsonData, err := json.Marshal(requestBody)
	if err != nil {
		http.Error(w, "Failed to marshal request body", http.StatusInternalServerError)
		return
	}

	// 創(chuàng)建 HTTP 請(qǐng)求
	req, err := http.NewRequest("POST", apiURL, bytes.NewBuffer(jsonData))
	if err != nil {
		http.Error(w, "Failed to create request", http.StatusInternalServerError)
		return
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+authToken)

	// 設(shè)置 HTTP 客戶端
	client := &http.Client{Timeout: time.Second * 50}
	resp, err := client.Do(req)
	if err != nil {
		http.Error(w, "Failed to get response", http.StatusInternalServerError)
		return
	}
	defer resp.Body.Close()

	// 設(shè)置響應(yīng)頭，開(kāi)啟流式輸出
	w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")

	// 確保 ResponseWriter 支持 Flusher
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
		return
	}

	// 處理流式響應(yīng)
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := scanner.Text()

		// 處理以 "data: " 開(kāi)頭的行
		if strings.HasPrefix(line, "data: ") {
			line = strings.TrimPrefix(line, "data: ")
		}
		if line == "[DONE]" {
			break
		}
		if line == "" {
			continue
		}

		// 解析響應(yīng)內(nèi)容
		var chunk ResponseBody
		if err := json.Unmarshal([]byte(line), &chunk); err != nil {
			continue
		}

		// 將響應(yīng)數(shù)據(jù)逐步發(fā)送給客戶端
		for _, choice := range chunk.Choices {
			content := choice.Delta.Content
			_, err := w.Write([]byte(content))
			if err != nil {
				http.Error(w, "Failed to write response", http.StatusInternalServerError)
				return
			}
			flusher.Flush() // 刷新緩沖區(qū)
		}
	}

	if err := scanner.Err(); err != nil {
		http.Error(w, "Scanner error", http.StatusInternalServerError)
		return
	}
}

func main() {
	http.HandleFunc("/stream", StreamHandler)
	fmt.Println("Server started at :8080")
	http.ListenAndServe(":8080", nil)
}

w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Connection", "keep-alive")