package wav2Lip

import (
	"context"
	"strconv"
)

type CheckPoint string

const (
	// CheckPointPathWav2Lip 高精度唇形同步
	CheckPointPathWav2Lip CheckPoint = "checkpoints/wav2lip.pth"
	// CheckPointPathWav2LipGAN 口型同步稍差，但视觉质量更好
	CheckPointPathWav2LipGAN CheckPoint = "checkpoints/wav2lip_gan.pth"
)

type Exec struct {
	pyExec      Python
	wav2lipPath string
}

type Python interface {
	ExecCtx(ctx context.Context, codeDir, pyFile string, args ...string) error
}

func New(p Python, wav2lipPath string) *Exec {
	return &Exec{
		pyExec:      p,
		wav2lipPath: wav2lipPath,
	}
}

// Inference 生成口型合成视频
// checkPointPath 用于加载权重的已保存检查点的名称
// face 包含要使用的面孔的视频/图像的文件路径
// audio 用作原始音频源的视频/音频文件的文件路径
// outFile 保存结果的视频路径
func (e *Exec) Inference(ctx context.Context, checkPointPath CheckPoint, face, audio, outFile string, options ...[]string) error {
	const py = "inference"
	var args []string
	args = append(args, optionCheckPointPath(string(checkPointPath))...)
	args = append(args, optionFace(face)...)
	args = append(args, optionAudio(audio)...)
	args = append(args, optionOutFile(outFile)...)
	for _, option := range options {
		args = append(args, option...)
	}
	return e.pyExec.ExecCtx(ctx, e.wav2lipPath, py, args...)
}

func optionCheckPointPath(val string) []string {
	return []string{
		"--checkpoint_path",
		val,
	}
}

func optionFace(val string) []string {
	return []string{
		"--face",
		val,
	}
}

func optionAudio(val string) []string {
	return []string{
		"--audio",
		val,
	}
}

func optionOutFile(val string) []string {
	return []string{
		"--outfile",
		val,
	}
}

// OptionPads 内边距（上、下、左、右）。 请调整至至少包括下巴
func OptionPads(top, bottom, left, right uint) []string {
	return []string{
		"--pads",
		strconv.Itoa(int(top)),
		strconv.Itoa(int(bottom)),
		strconv.Itoa(int(left)),
		strconv.Itoa(int(right)),
	}
}

// OptionResizeFactor 按此系数降低分辨率。
func OptionResizeFactor(val uint) []string {
	return []string{
		"--resize_factor",
		strconv.Itoa(int(val)),
	}
}

// OptionNoSmooth 如果您看到嘴部位置错位或一些奇怪的伪影（例如两张嘴），则可能是因为面部检测过度平滑。 使用 --nosmooth 参数并再试一次。
func OptionNoSmooth() []string {
	return []string{"--nosmooth"}
}