webp/bin/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"regexp"
	"strconv"

	"encoding/json"

	"git.satori.love/gameui/webp/api"
	"git.satori.love/gameui/webp/models"
	_ "github.com/go-sql-driver/mysql"
	"github.com/graphql-go/graphql"
	"github.com/graphql-go/handler"
	"github.com/spf13/viper"
)

// string 转换为 int, 如果转换失败则返回默认值
func stringToInt(str string, defaultValue int) int {
	if str == "" {
		return defaultValue
	}
	value, err := strconv.Atoi(str)
	if err != nil {
		return defaultValue
	}
	return value
}

func LogComponent(startTime int64, r *http.Request) {
	ms := (time.Now().UnixNano() - startTime) / 1000000
	color := "\033[1;32m%d\033[0m"
	if ms > 800 {
		color = "\033[1;31m%dms\033[0m" // 紅色加重
	} else if ms > 500 {
		color = "\033[1;33m%dms\033[0m" // 黃色加重
	} else if ms > 300 {
		color = "\033[1;32m%dms\033[0m" // 綠色加重
	} else if ms > 200 {
		color = "\033[1;34m%dms\033[0m" // 藍色加重
	} else if ms > 100 {
		color = "\033[1;35m%dms\033[0m" // 紫色加重
	} else {
		color = "\033[1;36m%dms\033[0m" // 黑色加重
	}
	endTime := fmt.Sprintf(color, ms)
	method := fmt.Sprintf("\033[1;32m%s\033[0m", r.Method) // 綠色加重
	url := fmt.Sprintf("\033[1;34m%s\033[0m", r.URL)       // 藍色加重

	log.Println(method, url, endTime, r.Header.Get("X-Forwarded-For"))
}

func LogRequest(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志

		var user_id int
		if token := r.Header.Get("token"); token != "" {
			user_id = api.ParseToken(token)
		}
		next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), "user_id", user_id)))
	})
}

var mysqlConnection models.MysqlConnection
var milvusConnection models.MilvusConnection

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU() - 1)

	configFilePath := flag.String("config", "./data/config.yaml", "配置文件路径")
	flag.Parse()

	viper.SetConfigFile(*configFilePath)
	if err := viper.ReadInConfig(); err != nil {
		log.Println("读取配置文件失败", err)
	}
	config := viper.GetViper()

	models.InitConfig(config)
	models.ZincInit()
	api.InitDefault(config)

	mysqlConnection.Init()
	milvusConnection.Init()
	err := milvusConnection.Client.LoadCollection(context.Background(), "default", false)
	if err != nil {
		log.Println("Milvus load collection failed:", err)
		return
	}

	if config.GetBool("oss.local") {
		fmt.Println("开启图像色调计算")
		go api.CheckColorNullRows(0)
	}
	if config.GetBool("gorse.open") {
		fmt.Println("开启用户偏好收集")
		api.GorseInit(config.GetString("gorse.host"), config.GetInt("gorse.port"))
		fmt.Println("同步点赞数据")
		go api.PutPraises()
	}

	schema, err := graphql.NewSchema(graphql.SchemaConfig{Query: graphql.NewObject(graphql.ObjectConfig{Name: "RootQuery", Fields: graphql.Fields{
		"users":       api.UserItems,
		"games":       api.GameItems,
		"works":       api.WorkItems,
		"collections": api.CollectionItems,
		"articles":    api.ArticleItems,
		"article":     api.ArticleItem,
		"images":      api.ImageItems,
		"image":       api.ImageItem,
		"searchs":     api.SearchItems,
	}})})

	if err != nil {
		log.Fatalf("failed to create new schema, error: %v", err)
	}

	http.Handle("/api", LogRequest(handler.New(&handler.Config{
		Schema:     &schema,
		Playground: true,
		Pretty:     false,
	})))

	// URL 格式: /image/{type}-{id}-{width}x{height}-{fit}.{format}?version
	http.HandleFunc("/image", func(w http.ResponseWriter, r *http.Request) {
		defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志

		// 如果本地文件存在，直接输出
		filePath := filepath.Join("data/webp", r.URL.Path)
		if _, err := os.Stat(filePath); err == nil {
			http.ServeFile(w, r, filePath)
			return
		}

		reg := regexp.MustCompile(`^/image/([a-z]+)-([0-9]+)-([0-9]+)x([0-9]+)-([a-z]+).(jpg|jpeg|png|webp)$`)
		matches := reg.FindStringSubmatch(r.URL.Path)
		if len(matches) != 7 {
			log.Println("URL 格式错误", matches)
			w.WriteHeader(http.StatusNotFound)
			return
		}
		group, id, width, height, fit, format := matches[1], matches[2], stringToInt(matches[3], 0), stringToInt(matches[4], 0), matches[5], matches[6]
		content, err := mysqlConnection.GetImageContent(group, id)
		if err != nil {
			log.Println("获取图片失败", format, err)
			w.WriteHeader(http.StatusNotFound)
			return
		}
		var img models.Image
		if err := img.Init(content); err != nil {
			log.Println("初始化图片失败", format, err)
			w.WriteHeader(http.StatusNotFound)
			return
		}
		data, err := img.ToWebP(width, height, fit)
		if err != nil {
			log.Println("转换图片失败", err)
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		err = os.MkdirAll(filepath.Dir(filePath), os.ModePerm)
		if err != nil {
			log.Println("创建文件目录失败:", err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		err = os.WriteFile(filePath, data, 0644)
		if err != nil {
			log.Println("保存文件失败:", err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		w.Header().Set("Content-Type", "image/webp")
		w.Header().Set("Cache-Control", "max-age=604800")
		w.Write(data)
	})

	// 获取转换后的m3u8视频链接
	http.HandleFunc("/video", func(w http.ResponseWriter, r *http.Request) {
		defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志

		queryParam := r.URL.Query().Get("url")
		safeParam, err := url.QueryUnescape(queryParam)
		if err != nil {
			log.Println("解码URL失败", err)
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		fmt.Println("safeParam", safeParam)
		urls, err := models.GetVideoM3U8(safeParam)
		fmt.Println("urls", urls, err)
		if err != nil {
			log.Println("获取视频链接失败", err)
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		// 将对象转换为有缩进的JSON输出
		json, _ := json.MarshalIndent(urls, "", "    ")
		w.Header().Set("Content-Type", "application/json")
		w.Write(json)
	})

	// 从Viper中读取配置
	port := viper.GetString("server.port")
	log.Println("Server is running at http://localhost:" + port)
	http.ListenAndServe(":"+port, nil)
}