package main

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"image"
	"io"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"

	"github.com/fogleman/gg"
	_ "golang.org/x/image/webp"
)

var maxFileSize = 200 << 20
var crfLevel = 24

func main() {
	_maxFileSize := os.Getenv("MAX_FILE_SIZE")
	if _maxFileSize != "" {
		size, err := strconv.Atoi(_maxFileSize)
		if err == nil {
			maxFileSize = size << 20
		}
	}
	log.Printf("Max file size: %d MB\n", maxFileSize/(1024*1024))

	_crfLevel := os.Getenv("CRF_LEVEL")
	if _crfLevel != "" {
		crfLevel, _ = strconv.Atoi(_crfLevel)
	}
	log.Printf("CRF level: %d\n", crfLevel)

	http.HandleFunc("/add-watermark", addWatermarkHandler)
	http.HandleFunc("/add-watermark-v2", addWatermarkHandlerV2)

	log.Println("Server starting on port 8080...")
	log.Fatal(http.ListenAndServe(":8080", nil))
}

func addWatermarkHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Only POST method is allowed", http.StatusMethodNotAllowed)
		return
	}
	crf := crfLevel
	_crf := r.URL.Query().Get("crf")
	if _crf != "" {
		crf2, err := strconv.Atoi(_crf)
		if err == nil {
			crf = crf2
			log.Printf("custom CRF level: %d\n", crf)
		}
	}

	// 解析multipart form
	err := r.ParseMultipartForm(10 << 20) // 10 MB
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// 获取上传的视频文件
	videoFile, videoHeader, err := r.FormFile("input")
	if err != nil {
		http.Error(w, "Error retrieving the video file", http.StatusBadRequest)
		return
	}
	defer videoFile.Close()
	// 获取上传的水印图片
	watermarkFile, watermarkHeader, err := r.FormFile("watermark")
	if err != nil {
		http.Error(w, "Error retrieving the watermark file", http.StatusBadRequest)
		return
	}

	defer watermarkFile.Close()
	watermarkPath := fmt.Sprintf("/%s_%s", randomId(), watermarkHeader.Filename)
	saveFile(watermarkFile, watermarkPath)
	defer os.RemoveAll(watermarkPath)

	// 创建临时文件
	videoPath := fmt.Sprintf("/%s_%s", randomId(), videoHeader.Filename)
	outputPath := fmt.Sprintf("/%s_%s", randomId(), "output_"+videoHeader.Filename)
	defer os.RemoveAll(videoPath)
	defer os.RemoveAll(outputPath)
	saveFile(videoFile, videoPath)

	start := time.Now()
	log.Printf("start processing video: %s\n", videoPath)
	cmd := exec.Command("ffmpeg",
		"-i", videoPath,
		"-i", watermarkPath,
		"-c:v", "libx264",
		"-c:a", "copy",
		"-crf", strconv.Itoa(crf),
		"-filter_complex", "[0:v]scale='min(1920\\,iw)':'-2'[scaled];[scaled][1:v]overlay=(main_w-overlay_w)/2:(main_h-overlay_h)-15[out]",
		"-map", "[out]",
		"-map", "0:a",
		outputPath)

	// 如果是图片
	imageSuffix := []string{"webp", ".png", ".jpg", ".jpeg", ".gif"}
	for _, suffix := range imageSuffix {
		if strings.HasSuffix(strings.ToLower(videoHeader.Filename), suffix) {
			cmd = exec.Command("ffmpeg",
				"-i", videoPath,
				"-i", watermarkPath,
				"-filter_complex", "overlay=(main_w-overlay_w)/2:(main_h-overlay_h)-15",
				outputPath)
			break
		}
	}

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	log.Printf("FFmpeg command: %s\n", cmd.String())
	err = cmd.Run()
	if err != nil {
		http.Error(w, "Error executing FFmpeg command: "+err.Error(), http.StatusInternalServerError)
		log.Printf("Error executing FFmpeg command: %s, %s\n", stderr.String(), stdout.String())
		return
	}

	log.Printf("finish processing video: %s, cost: %s\n", videoPath, time.Since(start))

	// 发送处理后的视频
	w.Header().Set("Content-Disposition", "attachment; filename=watermarked_"+videoHeader.Filename)
	w.Header().Set("Content-Type", videoHeader.Header.Get("Content-Type"))

	http.ServeFile(w, r, outputPath)
}

func addWatermarkHandlerV2(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Only POST method is allowed", http.StatusMethodNotAllowed)
		return
	}
	crf := crfLevel
	_crf := r.URL.Query().Get("crf")
	if _crf != "" {
		crf2, err := strconv.Atoi(_crf)
		if err == nil {
			crf = crf2
			log.Printf("custom CRF level: %d\n", crf)
		}
	}

	// 解析multipart form
	err := r.ParseMultipartForm(10 << 20) // 10 MB
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// 获取上传的视频文件
	videoFile, videoHeader, err := r.FormFile("input")
	if err != nil {
		http.Error(w, "Error retrieving the video file", http.StatusBadRequest)
		return
	}
	defer videoFile.Close()

	watermarkPaths := make([]string, 0)
	watermarkFiles := make([]multipart.File, 0)
	watermarkPositions := make([]string, 0)
	contextId := randomId()

	hasWatermark := false
	for i := 1; i <= 5; i++ {
		key := fmt.Sprintf("watermark_%d", i)
		watermarkFile, watermarkHeader, err := r.FormFile(key)
		if err != nil && err != http.ErrMissingFile {
			http.Error(w, "Error retrieving the watermark file", http.StatusBadRequest)
			return
		}
		if err == http.ErrMissingFile && !hasWatermark {
			http.Error(w, "Error retrieving the watermark file", http.StatusBadRequest)
			return
		}
		if err == http.ErrMissingFile && hasWatermark {
			break
		}
		hasWatermark = true
		watermarkPath := fmt.Sprintf("./%s_%s", contextId, watermarkHeader.Filename)
		saveFile(watermarkFile, watermarkPath)
		watermarkPaths = append(watermarkPaths, watermarkPath)
		watermarkFiles = append(watermarkFiles, watermarkFile)
		position := r.URL.Query().Get(key)
		if position == "" {
			http.Error(w, "Error retrieving the watermark position", http.StatusBadRequest)
			return
		}
		watermarkPositions = append(watermarkPositions, position)
	}

	defer func() {
		for _, f := range watermarkFiles {
			f.Close()
		}
		for _, p := range watermarkPaths {
			os.RemoveAll(p)
		}
	}()

	// 创建临时文件
	videoPath := fmt.Sprintf("./%s_%s", contextId, videoHeader.Filename)
	outputPath := fmt.Sprintf("./%s_%s", contextId, "output_"+videoHeader.Filename)
	defer os.RemoveAll(videoPath)
	defer os.RemoveAll(outputPath)
	saveFile(videoFile, videoPath)

	start := time.Now()
	log.Printf("start processing video: %s\n", videoPath)

	args := []string{"-i", videoPath}
	for i := 0; i < len(watermarkPaths); i++ {
		args = append(args, "-i", watermarkPaths[i])
	}

	args = append(args, "-c:v", "libx264", "-c:a", "copy", "-crf", strconv.Itoa(crf), "-filter_complex")

	filterComplex := fmt.Sprintf("[0:v]scale='min(1920\\,iw)':'-2'[tmp0];")
	for i := 0; i < len(watermarkPositions); i++ {
		filterComplex += fmt.Sprintf("[tmp%d][%d:v]overlay=%s[tmp%d];", i, i+1, watermarkPositions[i], i+1)
	}
	args = append(args, filterComplex[:len(filterComplex)-1], "-map", fmt.Sprintf("[tmp%d]", len(watermarkPaths)), "-map", "0:a", outputPath)
	// cmd := exec.Command("ffmpeg",
	//      "-i", videoPath,
	//      "-i", watermarkPath,
	//      "-c:v", "libx264",
	//      "-c:a", "copy",
	//      "-crf", strconv.Itoa(crf),
	//      "-filter_complex", "[0:v]scale='min(1920\\,iw)':'-2'[scaled];[scaled][1:v]overlay=(main_w-overlay_w)/2:(main_h-overlay_h)-15[out]",
	//      "-map", "[out]",
	//      "-map", "0:a",
	//      outputPath)

	cmd := exec.Command("ffmpeg", args...)

	// 如果是图片
	imageSuffix := []string{"webp", ".png", ".jpg", ".jpeg", ".gif"}
	for _, suffix := range imageSuffix {
		if strings.HasSuffix(strings.ToLower(videoHeader.Filename), suffix) {
			// 处理文字水印
			textArgs := r.FormValue("args")
			temp := struct {
				Text     string `json:"text"`
				Chinese  bool   `json:"chinese"`
				Fontsize int    `json:"fontsize"`
			}{}
			var overHeight int
			if textArgs != "" {
				err := json.Unmarshal([]byte(textArgs), &temp)
				if err != nil {
					log.Println("err1", err)
					http.Error(w, "Error retrieving the watermark text", http.StatusBadRequest)
					return
				}

				f, err := os.Open(videoPath)
				if err != nil {
					http.Error(w, "Error retrieving the watermark text", http.StatusBadRequest)
					return
				}

				img, _, err := image.Decode(f)
				if err != nil {
					log.Println("err2", err)
					http.Error(w, "Error retrieving the watermark text", http.StatusBadRequest)
					return
				}
				var ggContext *gg.Context
				ggContext, overHeight, err = drawText(img, temp.Text, temp.Fontsize, temp.Chinese)
				if err != nil {
					fmt.Println(err)
					http.Error(w, "Error retrieving the watermark text", http.StatusBadRequest)
					return
				}

				// 重新保存该图片
				err = ggContext.SavePNG(videoPath)
				if err != nil {
					log.Println("err3", err)
					http.Error(w, "Error retrieving the watermark text", http.StatusBadRequest)
					return
				}
			}

			args := []string{"-i", videoPath}
			for i := 0; i < len(watermarkPaths); i++ {
				args = append(args, "-i", watermarkPaths[i])
			}
			args = append(args, "-filter_complex")
			position := watermarkPositions[0]
			if overHeight > 0 && strings.Contains(position, "main_h") {
				position += "-" + strconv.Itoa(overHeight)
			}
			filterComplex := fmt.Sprintf("[0][1]overlay=%s[tmp0];", position)
			for i := 0; i < len(watermarkPositions[1:]); i++ {
				position = watermarkPositions[i+1]
				if overHeight > 0 && strings.Contains(position, "main_h") {
					position += "-" + strconv.Itoa(overHeight)
				}
				filterComplex += fmt.Sprintf("[tmp%d][%d]overlay=%s[tmp%d];", i, i+2, position, i+1)
			}

			args = append(args, filterComplex[:len(filterComplex)-1], "-map", fmt.Sprintf("[tmp%d]", len(watermarkPaths)-1), outputPath)

			// cmd = exec.Command("ffmpeg",
			//      "-i", videoPath,
			//      "-i", watermarkPath,
			//      "-filter_complex", "overlay=(main_w-overlay_w)/2:(main_h-overlay_h)-15",
			//      outputPath)
			cmd = exec.Command("ffmpeg", args...)
			break
		}
	}

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	log.Printf("FFmpeg command: %s\n", cmd.String())

	err = cmd.Run()
	if err != nil {
		http.Error(w, "Error executing FFmpeg command: "+err.Error(), http.StatusInternalServerError)
		log.Printf("Error executing FFmpeg command: %s, %s\n", stderr.String(), stdout.String())
		return
	}

	log.Printf("finish processing video: %s, cost: %s\n", videoPath, time.Since(start))

	// 发送处理后的视频
	w.Header().Set("Content-Disposition", "attachment; filename=watermarked_"+videoHeader.Filename)
	w.Header().Set("Content-Type", videoHeader.Header.Get("Content-Type"))

	http.ServeFile(w, r, outputPath)
}

func saveFile(file io.Reader, path string) error {
	out, err := os.Create(path)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, file)
	return err
}

func randomId() string {
	b := make([]byte, 16)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}
