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/milvus-io/milvus-sdk-go/v2/entity"
	"github.com/spf13/viper"

	lru "github.com/hashicorp/golang-lru/v2"
)

// 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)))
	})
}

type Image struct {
	Id                   int            `json:"id"                      db:"id"`
	Width                int            `json:"width"                   db:"width"`
	Height               int            `json:"height"                  db:"height"`
	Content              string         `json:"content"                 db:"content"`
	ArticleCategoryTopId int            `json:"article_category_top_id" db:"article_category_top_id"`
	PraiseCount          int            `json:"praise_count"            db:"praise_count"`
	CollectCount         int            `json:"collect_count"           db:"collect_count"`
	CreateTime           time.Time      `json:"createTime"              db:"createTime"`
	UpdateTime           time.Time      `json:"updateTime"              db:"updateTime"`
	UserID               int            `json:"user_id"                 db:"user_id"`
	User                 models.User    `json:"user"                    db:"user"`
	Article              models.Article `json:"article"                 db:"article"`
}

type Tag struct {
	Id         int       `json:"id"`
	Name       string    `json:"name"`
	CreateTime time.Time `json:"create_time"`
	UpdateTime time.Time `json:"update_time"`
}

type History struct {
	Type       string      `json:"type"`
	CreateTime time.Time   `json:"create_time"`
	Data       interface{} `json:"data"`
}

type ListView struct {
	Code     int           `json:"code"`
	Page     int           `json:"page"`
	PageSize int           `json:"pageSize"`
	Total    int           `json:"total"`
	Next     bool          `json:"next"`
	List     []interface{} `json:"list"`
}

var mysqlConnection models.MysqlConnection
var milvusConnection models.MilvusConnection

func GetNetWorkEmbedding(id int) (embedding []float32) {
	host := viper.GetString("embedding.host")
	port := viper.GetInt("embedding.port")
	httpClient := &http.Client{}
	request, err := http.NewRequest("PUT", fmt.Sprintf("http://%s:%d/api/default/%d", host, port, id), nil)
	if err != nil {
		log.Println("请求失败1:", err)
		return
	}
	response, err := httpClient.Do(request)
	if err != nil {
		log.Println("请求失败2:", err)
		return
	}
	defer response.Body.Close()

	var result struct {
		Code    int       `json:"code"`
		Message string    `json:"message"`
		Feature []float32 `json:"feature"`
	}
	err = json.NewDecoder(response.Body).Decode(&result)
	if err != nil {
		log.Println("解析失败:", err)
		return
	}
	if result.Code != 0 {
		log.Println("请求失败3:", result.Message)
		return
	}
	return result.Feature
}

var lruCache, _ = lru.New[int, []int64](100000)

func (image *Image) GetSimilarImagesIdList(collection_name string) (ids []int64) {
	ctx := context.Background()

	// 先从 LRU 中查询缓存的结果, 如果缓存中有, 直接返回
	if value, ok := lruCache.Get(image.Id); ok {
		return value
	}

	// 先从milvus中查询图片的向量
	var embedding []float32
	result, err := milvusConnection.Client.Query(ctx, collection_name, nil, fmt.Sprintf("id in [%d]", image.Id), []string{"embedding"})
	if err != nil {
		log.Println("查詢向量失敗:", err)
		embedding = GetNetWorkEmbedding(image.Id)
	} else {
		for _, item := range result {
			if item.Name() == "embedding" {
				embedding = item.FieldData().GetVectors().GetFloatVector().Data
				continue
			}
		}
	}

	// 处理向量不存在的情况
	if len(embedding) == 0 {
		log.Println("向量不存在, 也未能重新生成")
		return ids
	}

	// 用向量查询相似图片
	topk := 200
	sp, _ := entity.NewIndexIvfFlatSearchParam(64)
	vectors := []entity.Vector{entity.FloatVector(embedding)}
	resultx, err := milvusConnection.Client.Search(ctx, collection_name, nil, "", []string{"id", "article_id"}, vectors, "embedding", entity.L2, topk, sp)
	if err != nil {
		log.Println("搜索相似失敗:", err)
		return
	}

	// 输出结果
	for _, item := range resultx {
		ids = item.IDs.FieldData().GetScalars().GetLongData().GetData()
	}

	// 将结果缓存到 LRU 中
	lruCache.Add(image.Id, ids)

	return ids
}

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"))
	}

	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)
}