main.go

package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/joho/godotenv"
)

type Config struct {
	APIPrefix     string
	APIKey        string
	MaxRetryCount int
	RetryDelay    time.Duration
	FakeHeaders   map[string]string
}

var config Config

func init() {
	godotenv.Load()
	config = Config{
		APIKey:        getEnv("API_KEY", ""),
		MaxRetryCount: getIntEnv("MAX_RETRY_COUNT", 3),
		RetryDelay:    getDurationEnv("RETRY_DELAY", 5000),
		FakeHeaders: map[string]string{
			"Accept":             "*/*",
			"Accept-Encoding":    "gzip, deflate, br, zstd",
			"Accept-Language":    "zh-CN,zh;q=0.9",
			"Origin":             "https://duckduckgo.com/",
			"Cookie":             "l=wt-wt; ah=wt-wt; dcm=6",
			"Dnt":                "1",
			"Priority":           "u=1, i",
			"Referer":            "https://duckduckgo.com/",
			"Sec-Ch-Ua":          `"Microsoft Edge";v="129", "Not(A:Brand";v="8", "Chromium";v="129"`,
			"Sec-Ch-Ua-Mobile":   "?0",
			"Sec-Ch-Ua-Platform": `"Windows"`,
			"Sec-Fetch-Dest":     "empty",
			"Sec-Fetch-Mode":     "cors",
			"Sec-Fetch-Site":     "same-origin",
			"User-Agent":         "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36",
		},
	}
}

func authMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		apiKey := c.GetHeader("Authorization")
		if apiKey == "" {
			apiKey = c.Query("api_key")
		}

		// 当未提供或者 API Key 不合法时，允许匿名访问
		if apiKey == "" || !strings.HasPrefix(apiKey, "Bearer ") {
			c.Next() // 未提供 API 密钥时允许继续请求
			return
		}

		apiKey = strings.TrimPrefix(apiKey, "Bearer ")

		if apiKey != config.APIKey {
			c.JSON(http.StatusUnauthorized, gin.H{"error": "Invalid API key"})
			c.Abort()
			return
		}

		c.Next()
	}
}

func main() {
	gin.SetMode(gin.ReleaseMode)
	r := gin.Default()
	r.Use(corsMiddleware())

	// 健康检查
	r.GET("/ping", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{"message": "pang"})
	})

	// API 路由组
	apiGroup := r.Group("/")
	apiGroup.Use(authMiddleware()) // 可以选择性地提供 API 密钥
	{
		// 原始路径 /hf/v1/*
		apiGroup.GET("/hf/v1/models", handleModels)
		apiGroup.POST("/hf/v1/chat/completions", handleCompletion)

		// 新路径 /api/v1/*
		apiGroup.GET("/api/v1/models", handleModels)
		apiGroup.POST("/api/v1/chat/completions", handleCompletion)

		// 新路径 /v1/*
		apiGroup.GET("/v1/models", handleModels)
		apiGroup.POST("/v1/chat/completions", handleCompletion)

		// 新路径 /completions
		apiGroup.POST("/completions", handleCompletion)
	}

	// 从环境变量中读取端口号
	port := os.Getenv("PORT")
	if port == "" {
		port = "7860"
	}
	r.Run(":" + port)
}

func handleModels(c *gin.Context) {
	models := []gin.H{
		{"id": "gpt-4o-mini", "object": "model", "owned_by": "ddg"},
		{"id": "claude-3-haiku", "object": "model", "owned_by": "ddg"},
		{"id": "llama-3.1-70b", "object": "model", "owned_by": "ddg"},
		{"id": "mixtral-8x7b", "object": "model", "owned_by": "ddg"},
	}
	c.JSON(http.StatusOK, gin.H{"object": "list", "data": models})
}

func handleCompletion(c *gin.Context) {
	var req struct {
		Model    string `json:"model"`
		Messages []struct {
			Role    string      `json:"role"`
			Content interface{} `json:"content"`
		} `json:"messages"`
		Stream bool `json:"stream"`
	}

	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	model := convertModel(req.Model)
	content := prepareMessages(req.Messages)

	reqBody := map[string]interface{}{
		"model": model,
		"messages": []map[string]interface{}{
			{
				"role":    "user",
				"content": content,
			},
		},
	}

	body, err := json.Marshal(reqBody)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("请求体序列化失败: %v", err)})
		return
	}

	token, err := requestToken()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "无法获取token"})
		return
	}

	upstreamReq, err := http.NewRequest("POST", "https://duckduckgo.com/duckchat/v1/chat", strings.NewReader(string(body)))
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("创建请求失败: %v", err)})
		return
	}

	for k, v := range config.FakeHeaders {
		upstreamReq.Header.Set(k, v)
	}
	upstreamReq.Header.Set("x-vqd-4", token)
	upstreamReq.Header.Set("Content-Type", "application/json")
	//// just for print info
	//jsonData, err := json.Marshal(upstreamReq.Header)
	//if err == nil {
	//	log.Printf("Request Header: %s", jsonData)
	//}
	//log.Printf("Request Body: %s", body)
	client := &http.Client{
		Timeout: 30 * time.Second,
	}
	resp, err := client.Do(upstreamReq)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("请求失败: %v", err)})
		return
	}
	defer resp.Body.Close()

	if req.Stream {
		// 启用 SSE 流式响应
		c.Writer.Header().Set("Content-Type", "text/event-stream")
		c.Writer.Header().Set("Cache-Control", "no-cache")
		c.Writer.Header().Set("Connection", "keep-alive")

		flusher, ok := c.Writer.(http.Flusher)
		if !ok {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Streaming not supported"})
			return
		}

		reader := bufio.NewReader(resp.Body)
		for {
			line, err := reader.ReadString('\n')
			if err != nil {
				if err != io.EOF {
					log.Printf("读取流式响应失败: %v", err)
				}
				break
			}
			if strings.HasPrefix(line, "data: ") {
				// 解析响应中的 JSON 数据块
				line = strings.TrimPrefix(line, "data: ")
				line = strings.TrimSpace(line)
				// 忽略非 JSON 数据块（例如特殊标记 [DONE]）
				if line == "[DONE]" {
					break
				}
				var chunk map[string]interface{}
				if err := json.Unmarshal([]byte(line), &chunk); err != nil {
					log.Printf("解析响应行失败: %v", err)
					continue
				}

				// 检查 chunk 是否包含 message
				if msg, exists := chunk["message"]; exists && msg != nil {
					if msgStr, ok := msg.(string); ok {
						response := map[string]interface{}{
							"id":      "chatcmpl-QXlha2FBbmROaXhpZUFyZUF3ZXNvbWUK",
							"object":  "chat.completion.chunk",
							"created": time.Now().Unix(),
							"model":   model,
							"choices": []map[string]interface{}{
								{
									"index": 0,
									"delta": map[string]string{
										"content": msgStr,
									},
									"finish_reason": nil,
								},
							},
						}
						// 将响应格式化为 SSE 数据块
						sseData, _ := json.Marshal(response)
						sseMessage := fmt.Sprintf("data: %s\n\n", sseData)

						// 发送数据并刷新缓冲区
						_, writeErr := c.Writer.Write([]byte(sseMessage))
						if writeErr != nil {
							log.Printf("写入响应失败: %v", writeErr)
							break
						}
						flusher.Flush()
					} else {
						log.Printf("chunk[message] 不是字符串: %v", msg)
					}
				} else {
					// 解析行中有空行
					log.Println("chunk 中未包含 message 或 message 为 nil")
				}
			}
		}
	} else {
		// 非流式响应，返回完整的 JSON
		var fullResponse strings.Builder
		reader := bufio.NewReader(resp.Body)

		for {
			line, err := reader.ReadString('\n')
			if err == io.EOF {
				break
			} else if err != nil {
				log.Printf("读取响应失败: %v", err)
				break
			}

			if strings.HasPrefix(line, "data: ") {
				line = strings.TrimPrefix(line, "data: ")
				line = strings.TrimSpace(line)

				if line == "[DONE]" {
					break
				}

				var chunk map[string]interface{}
				if err := json.Unmarshal([]byte(line), &chunk); err != nil {
					log.Printf("解析响应行失败: %v", err)
					continue
				}

				if message, exists := chunk["message"]; exists {
					if msgStr, ok := message.(string); ok {
						fullResponse.WriteString(msgStr)
					}
				}
			}
		}

		// 返回完整 JSON 响应
		response := map[string]interface{}{
			"id":      "chatcmpl-QXlha2FBbmROaXhpZUFyZUF3ZXNvbWUK",
			"object":  "chat.completion",
			"created": time.Now().Unix(),
			"model":   model,
			"usage": map[string]int{
				"prompt_tokens":     0,
				"completion_tokens": 0,
				"total_tokens":      0,
			},
			"choices": []map[string]interface{}{
				{
					"message": map[string]string{
						"role":    "assistant",
						"content": fullResponse.String(),
					},
					"index": 0,
				},
			},
		}

		c.JSON(http.StatusOK, response)
	}
}

func requestToken() (string, error) {
	url := "https://duckduckgo.com/duckchat/v1/status"
	client := &http.Client{
		Timeout: 15 * time.Second, // 设置超时时间
	}

	maxRetries := config.MaxRetryCount
	retryDelay := config.RetryDelay

	for attempt := 0; attempt < maxRetries; attempt++ {
		if attempt > 0 {
			log.Printf("requestToken: 第 %d 次重试，等待 %v...", attempt, retryDelay)
			time.Sleep(retryDelay)
		}
		//log.Printf("requestToken: 发送 GET 请求到 %s", url)

		// 创建请求
		req, err := http.NewRequest("GET", url, nil)
		if err != nil {
			log.Printf("requestToken: 创建请求失败: %v", err)
			return "", fmt.Errorf("无法创建请求: %w", err)
		}

		// 添加假头部
		for k, v := range config.FakeHeaders {
			req.Header.Set(k, v)
		}
		req.Header.Set("x-vqd-accept", "1")

		// 发送请求
		resp, err := client.Do(req)
		if err != nil {
			log.Printf("requestToken: 请求失败: %v", err)
			continue // 网络通信失败，进行重试
		}
		defer resp.Body.Close()

		// 检查状态码是否为 200
		if resp.StatusCode != http.StatusOK {
			bodyBytes, _ := io.ReadAll(resp.Body) // 读取响应体，错误时也需要记录响应内容
			bodyString := string(bodyBytes)
			log.Printf("requestToken: 非200响应，状态码=%d, 响应内容: %s", resp.StatusCode, bodyString)
			continue
		}

		// 尝试从头部提取 token
		token := resp.Header.Get("x-vqd-4")
		if token == "" {
			log.Println("requestToken: 响应中未包含 x-vqd-4 头部")
			bodyBytes, _ := io.ReadAll(resp.Body)
			bodyString := string(bodyBytes)
			log.Printf("requestToken: 响应内容: %s", bodyString)
			continue
		}

		// 成功获取到 token
		//log.Printf("requestToken: 成功获取到 token: %s", token)
		return token, nil
	}

	// 如果所有重试均失败，返回错误
	return "", errors.New("requestToken: 无法获取到 token，多次重试仍失败")
}

func prepareMessages(messages []struct {
	Role    string      `json:"role"`
	Content interface{} `json:"content"`
}) string {
	var contentBuilder strings.Builder

	for _, msg := range messages {
		// Determine the role - 'system' becomes 'user'
		role := msg.Role
		if role == "system" {
			role = "user"
		}

		// Process the content as string
		contentStr := ""
		switch v := msg.Content.(type) {
		case string:
			contentStr = v
		case []interface{}:
			for _, item := range v {
				if itemMap, ok := item.(map[string]interface{}); ok {
					if text, exists := itemMap["text"].(string); exists {
						contentStr += text
					}
				}
			}
		default:
			contentStr = fmt.Sprintf("%v", msg.Content)
		}

		// Append the role and content to the builder
		contentBuilder.WriteString(fmt.Sprintf("%s:%s;\r\n", role, contentStr))
	}

	return contentBuilder.String()
}

func convertModel(inputModel string) string {
	switch strings.ToLower(inputModel) {
	case "claude-3-haiku":
		return "claude-3-haiku-20240307"
	case "llama-3.1-70b":
		return "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo"
	case "mixtral-8x7b":
		return "mistralai/Mixtral-8x7B-Instruct-v0.1"
	default:
		return "gpt-4o-mini"
	}
}

func corsMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
		c.Writer.Header().Set("Access-Control-Allow-Methods", "*")
		c.Writer.Header().Set("Access-Control-Allow-Headers", "*")
		if c.Request.Method == http.MethodOptions {
			c.AbortWithStatus(http.StatusNoContent)
			return
		}
		c.Next()
	}
}

func getEnv(key, fallback string) string {
	if value, exists := os.LookupEnv(key); exists {
		return value
	}
	return fallback
}

func getIntEnv(key string, fallback int) int {
	if value, exists := os.LookupEnv(key); exists {
		var intValue int
		fmt.Sscanf(value, "%d", &intValue)
		return intValue
	}
	return fallback
}

func getDurationEnv(key string, fallback int) time.Duration {
	return time.Duration(getIntEnv(key, fallback)) * time.Millisecond
}