package wav2lipEx

import (
	"context"
	"errors"
	"gitee.com/kaylee595/go-wav2lipEx/services/ffmpeg"
	"gitee.com/kaylee595/go-wav2lipEx/services/gfpgan"
	"gitee.com/kaylee595/go-wav2lipEx/services/wav2Lip"
	"gitee.com/kaylee595/go-wav2lipEx/utils"
	"go.uber.org/zap"
	"os"
	"path/filepath"
	"time"
)

type Wav2lipEx struct {
	wav2Lip Wav2Lip
	gfpgan  Gfpgan
	ffmpeg  *ffmpeg.Exec
	logger  *zap.Logger
}

type Wav2Lip interface {
	Inference(ctx context.Context, checkPointPath wav2Lip.CheckPoint, face, audio, outFile string, options ...[]string) error
}

type Gfpgan interface {
	InferenceGfpgan(ctx context.Context, input, output string, ver gfpgan.Version, options ...[]string) error
}

type Option func(*Wav2lipEx)

func New(options ...Option) (*Wav2lipEx, error) {
	w := &Wav2lipEx{}
	for _, option := range options {
		option(w)
	}
	if w.wav2Lip == nil {
		return nil, errors.New("wav2lip is required")
	}
	if w.gfpgan == nil {
		return nil, errors.New("gfpgan is required")
	}
	if w.ffmpeg == nil {
		return nil, errors.New("ffmpeg is required")
	}
	if w.logger == nil {
		w.logger = zap.L()
	}
	return w, nil
}

func WithWav2Lip(wav2Lip Wav2Lip) Option {
	return func(w *Wav2lipEx) {
		w.wav2Lip = wav2Lip
	}
}

func WithGfpgan(gfpgan Gfpgan) Option {
	return func(w *Wav2lipEx) {
		w.gfpgan = gfpgan
	}
}

func WithFfmpeg(ffmpeg *ffmpeg.Exec) Option {
	return func(w *Wav2lipEx) {
		w.ffmpeg = ffmpeg
	}
}

func WithLogger(logger *zap.Logger) Option {
	return func(w *Wav2lipEx) {
		w.logger = logger
	}
}

type Paras struct {
	// ModifyFace 更换脸部, 传入脸部照片/视频
	//ModifyFace string
	// UseEsrgan 默认使用gfpgan；你也可以将此参数设置TRUE，这将使用ESRGAN。gfpgan用于修复人脸，esrgan用于修复景色和动漫
	UseEsrgan bool
	// Nosmooth 防止在短时间窗口内平滑人脸检测
	Nosmooth bool
	// Pads 内边距（上、下、左、右）。 请调整至至少包括下巴
	Pads []uint
	// NoRepair 不修复照片
	NoRepair bool
}

// GenerateVideo
// 生成视频，传入一段嘴巴闭合状态的视频/图像和一段音频，将输出视频唇形与音频同步的画面到指定文件下.
// 会在输出目录下创建一个用于存放高清修复后的照片, 创建文件夹名字随机起名, 如果推理成功, 那么会自动清理该自动创建的文件夹, 如果推理失败
// 则会保留高清修复输出的图片
func (w *Wav2lipEx) GenerateVideo(ctx context.Context, input, audio, output string, pars Paras) (err error) {
	// 创建一个临时空间
	var tempDir string
	if !pars.NoRepair {
		tempDir, err = utils.CreateTempDir(filepath.Dir(output))
		if err != nil {
			w.logger.Error("创建临时文件夹失败", zap.Error(err))
			return err
		}
		w.logger.Info("创建临时文件夹", zap.String("tempDir", tempDir))
		defer func() {
			if err == nil {
				os.RemoveAll(tempDir)
			}
		}()
	}

	// 口型推理
	var wav2LipVideoFile string
	if pars.NoRepair {
		wav2LipVideoFile = output
	} else {
		wav2LipVideoFile = filepath.Join(tempDir, "wav2lipVideo.mp4")
	}
	w.logger.Info("口型推理", zap.String("wav2LipVideoFile", wav2LipVideoFile))
	var args [][]string
	args = append(args, wav2Lip.OptionResizeFactor(2))
	if pars.Nosmooth {
		args = append(args, wav2Lip.OptionNoSmooth())
	}
	if len(pars.Pads) >= 4 {
		args = append(args, wav2Lip.OptionPads(pars.Pads[0], pars.Pads[1], pars.Pads[2], pars.Pads[3]))
	}
	err = w.wav2Lip.Inference(ctx,
		wav2Lip.CheckPointPathWav2Lip,
		input,
		audio,
		wav2LipVideoFile,
		args...,
	)
	if err != nil {
		w.logger.Error("口型推理失败",
			zap.Error(err),
			zap.String("input", input),
			zap.String("audio", audio),
		)
		return err
	}

	if !pars.NoRepair {
		// 视频帧提取
		wav2LipVideoFramesDir := utils.JoinDir(tempDir, "wav2lipVideoFrames")
		w.logger.Info("视频帧提取", zap.String("wav2LipVideoFramesDir", wav2LipVideoFramesDir))
		err = w.ffmpeg.VideoFrameExtraction(ctx,
			wav2LipVideoFile,
			wav2LipVideoFramesDir,
		)
		if err != nil {
			w.logger.Error("视频帧提取失败", zap.Error(err), zap.String("videoFile", wav2LipVideoFile))
			return err
		}

		// 图像高清修复
		imageHDRepairFramesDir := utils.JoinDir(tempDir, "imageHDRepairFrames")
		w.logger.Info("图像高清修复", zap.String("imageHDRepairFramesDir", imageHDRepairFramesDir))
		c, cancelFunc := context.WithTimeout(ctx, time.Hour)
		defer cancelFunc()
		if pars.UseEsrgan {
			//err = esrgan.New(w.p, w.esrganPath).InferenceRealesrgan(c,
			//	wav2LipVideoFramesDir,
			//	imageHDRepairFramesDir,
			//	// 下面两个参数用于检测人脸后使用gfpgan修复
			//	esrgan.OptionFaceEnhance(),
			//	esrgan.OptionFp32(),
			//)
			//if err != nil {
			//	return err
			//}
		} else {
			err = w.gfpgan.InferenceGfpgan(c,
				wav2LipVideoFramesDir,
				imageHDRepairFramesDir,
				gfpgan.VersionV14,
				gfpgan.OptionBgTile(512),
				gfpgan.OptionUpScale(2),
				// 下面参数用于背景不使用esrgan修复
				//gfpgan.OptionBgUpSampler("none"),
			)
			if err != nil {
				w.logger.Error("图像高清修复失败", zap.Error(err), zap.String("wav2LipVideoFramesDir", wav2LipVideoFramesDir))
				return err
			}
		}

		// 图像合成到视频
		w.logger.Info("图像合成到视频")
		imageHDRepairFramesDir = filepath.Join(imageHDRepairFramesDir, "restored_imgs")
		err = w.ffmpeg.MergeImg2Video(ctx,
			imageHDRepairFramesDir,
			audio,
			output,
		)
		if err != nil {
			w.logger.Error("图像合成到视频失败",
				zap.String("imageHDRepairFramesDir", imageHDRepairFramesDir),
				zap.String("audio", audio),
				zap.String("output", output),
			)
			return err
		}
	}
	return
}