package manager

import (
	"context"
	"fmt"
	"time"

	"github.com/remeh/sizedwaitgroup"
	"github.com/stashapp/stash/internal/manager/config"
	"github.com/stashapp/stash/pkg/image"
	"github.com/stashapp/stash/pkg/job"
	"github.com/stashapp/stash/pkg/logger"
	"github.com/stashapp/stash/pkg/models"
	"github.com/stashapp/stash/pkg/scene"
	"github.com/stashapp/stash/pkg/scene/generate"
	"github.com/stashapp/stash/pkg/sliceutil/stringslice"
)

type GenerateMetadataInput struct {
	Covers              bool                         `json:"covers"`
	Sprites             bool                         `json:"sprites"`
	Previews            bool                         `json:"previews"`
	ImagePreviews       bool                         `json:"imagePreviews"`
	PreviewOptions      *GeneratePreviewOptionsInput `json:"previewOptions"`
	Markers             bool                         `json:"markers"`
	MarkerImagePreviews bool                         `json:"markerImagePreviews"`
	MarkerScreenshots   bool                         `json:"markerScreenshots"`
	Transcodes          bool                         `json:"transcodes"`
	// Generate transcodes even if not required
	ForceTranscodes           bool `json:"forceTranscodes"`
	Phashes                   bool `json:"phashes"`
	InteractiveHeatmapsSpeeds bool `json:"interactiveHeatmapsSpeeds"`
	ClipPreviews              bool `json:"clipPreviews"`
	ImageThumbnails           bool `json:"imageThumbnails"`
	// scene ids to generate for
	SceneIDs []string `json:"sceneIDs"`
	// marker ids to generate for
	MarkerIDs []string `json:"markerIDs"`
	// overwrite existing media
	Overwrite bool `json:"overwrite"`
}

type GeneratePreviewOptionsInput struct {
	// Number of segments in a preview file
	PreviewSegments *int `json:"previewSegments"`
	// Preview segment duration, in seconds
	PreviewSegmentDuration *float64 `json:"previewSegmentDuration"`
	// Duration of start of video to exclude when generating previews
	PreviewExcludeStart *string `json:"previewExcludeStart"`
	// Duration of end of video to exclude when generating previews
	PreviewExcludeEnd *string `json:"previewExcludeEnd"`
	// Preset when generating preview
	PreviewPreset *models.PreviewPreset `json:"previewPreset"`
}

const generateQueueSize = 200000

type GenerateJob struct {
	repository models.Repository
	input      GenerateMetadataInput

	overwrite      bool
	fileNamingAlgo models.HashAlgorithm

	totals totalsGenerate
}

type totalsGenerate struct {
	covers                   int64
	sprites                  int64
	previews                 int64
	imagePreviews            int64
	markers                  int64
	transcodes               int64
	phashes                  int64
	interactiveHeatmapSpeeds int64
	clipPreviews             int64
	imageThumbnails          int64

	tasks int
}

func (j *GenerateJob) Execute(ctx context.Context, progress *job.Progress) error {
	var scenes []*models.Scene
	var err error
	var markers []*models.SceneMarker

	j.overwrite = j.input.Overwrite
	j.fileNamingAlgo = config.GetInstance().GetVideoFileNamingAlgorithm()

	config := config.GetInstance()
	parallelTasks := config.GetParallelTasksWithAutoDetection()

	logger.Infof("Generate started with %d parallel tasks", parallelTasks)

	queue := make(chan Task, generateQueueSize)
	go func() {
		defer close(queue)

		sceneIDs, err := stringslice.StringSliceToIntSlice(j.input.SceneIDs)
		if err != nil {
			logger.Error(err.Error())
		}
		markerIDs, err := stringslice.StringSliceToIntSlice(j.input.MarkerIDs)
		if err != nil {
			logger.Error(err.Error())
		}

		g := &generate.Generator{
			Encoder:      instance.FFMpeg,
			FFMpegConfig: instance.Config,
			LockManager:  instance.ReadLockManager,
			MarkerPaths:  instance.Paths.SceneMarkers,
			ScenePaths:   instance.Paths.Scene,
			Overwrite:    j.overwrite,
		}

		r := j.repository
		if err := r.WithReadTxn(ctx, func(ctx context.Context) error {
			qb := r.Scene
			if len(j.input.SceneIDs) == 0 && len(j.input.MarkerIDs) == 0 {
				j.queueTasks(ctx, g, queue)
			} else {
				if len(j.input.SceneIDs) > 0 {
					scenes, err = qb.FindMany(ctx, sceneIDs)
					for _, s := range scenes {
						if err := s.LoadFiles(ctx, qb); err != nil {
							return err
						}

						j.queueSceneJobs(ctx, g, s, queue)
					}
				}

				if len(j.input.MarkerIDs) > 0 {
					markers, err = r.SceneMarker.FindMany(ctx, markerIDs)
					if err != nil {
						return err
					}
					for _, m := range markers {
						j.queueMarkerJob(g, m, queue)
					}
				}
			}

			return nil
		}); err != nil && ctx.Err() == nil {
			logger.Error(err.Error())
			return
		}

		totals := j.totals
		logMsg := "Generating"
		if j.input.Covers {
			logMsg += fmt.Sprintf(" %d covers", totals.covers)
		}
		if j.input.Sprites {
			logMsg += fmt.Sprintf(" %d sprites", totals.sprites)
		}
		if j.input.Previews {
			logMsg += fmt.Sprintf(" %d previews", totals.previews)
		}
		if j.input.ImagePreviews {
			logMsg += fmt.Sprintf(" %d image previews", totals.imagePreviews)
		}
		if j.input.Markers {
			logMsg += fmt.Sprintf(" %d markers", totals.markers)
		}
		if j.input.Transcodes {
			logMsg += fmt.Sprintf(" %d transcodes", totals.transcodes)
		}
		if j.input.Phashes {
			logMsg += fmt.Sprintf(" %d phashes", totals.phashes)
		}
		if j.input.InteractiveHeatmapsSpeeds {
			logMsg += fmt.Sprintf(" %d heatmaps & speeds", totals.interactiveHeatmapSpeeds)
		}
		if j.input.ClipPreviews {
			logMsg += fmt.Sprintf(" %d Image Clip Previews", totals.clipPreviews)
		}
		if j.input.ImageThumbnails {
			logMsg += fmt.Sprintf(" %d Image Thumbnails", totals.imageThumbnails)
		}
		if logMsg == "Generating" {
			logMsg = "Nothing selected to generate"
		}
		logger.Infof(logMsg)

		progress.SetTotal(int(totals.tasks))
	}()

	wg := sizedwaitgroup.New(parallelTasks)

	// Start measuring how long the generate has taken. (consider moving this up)
	start := time.Now()
	if err = instance.Paths.Generated.EnsureTmpDir(); err != nil {
		logger.Warnf("could not create temporary directory: %v", err)
	}

	defer func() {
		if err := instance.Paths.Generated.EmptyTmpDir(); err != nil {
			logger.Warnf("failure emptying temporary directory: %v", err)
		}
	}()

	for f := range queue {
		if job.IsCancelled(ctx) {
			break
		}

		wg.Add()
		// #1879 - need to make a copy of f - otherwise there is a race condition
		// where f is changed when the goroutine runs
		localTask := f
		go progress.ExecuteTask(localTask.GetDescription(), func() {
			localTask.Start(ctx)
			wg.Done()
			progress.Increment()
		})
	}

	wg.Wait()

	if job.IsCancelled(ctx) {
		logger.Info("Stopping due to user request")
		return nil
	}

	elapsed := time.Since(start)
	logger.Info(fmt.Sprintf("Generate finished (%s)", elapsed))
	return nil
}

func (j *GenerateJob) queueTasks(ctx context.Context, g *generate.Generator, queue chan<- Task) {
	j.totals = totalsGenerate{}

	j.queueScenesTasks(ctx, g, queue)
	j.queueImagesTasks(ctx, g, queue)
}

func (j *GenerateJob) queueScenesTasks(ctx context.Context, g *generate.Generator, queue chan<- Task) {
	const batchSize = 1000

	findFilter := models.BatchFindFilter(batchSize)

	r := j.repository

	for more := true; more; {
		if job.IsCancelled(ctx) {
			return
		}

		scenes, err := scene.Query(ctx, r.Scene, nil, findFilter)
		if err != nil {
			logger.Errorf("Error encountered queuing files to scan: %s", err.Error())
			return
		}

		for _, ss := range scenes {
			if job.IsCancelled(ctx) {
				return
			}

			if err := ss.LoadFiles(ctx, r.Scene); err != nil {
				logger.Errorf("Error encountered queuing files to scan: %s", err.Error())
				return
			}

			j.queueSceneJobs(ctx, g, ss, queue)
		}

		if len(scenes) != batchSize {
			more = false
		} else {
			*findFilter.Page++
		}
	}
}

func (j *GenerateJob) queueImagesTasks(ctx context.Context, g *generate.Generator, queue chan<- Task) {
	const batchSize = 1000

	findFilter := models.BatchFindFilter(batchSize)

	r := j.repository

	for more := j.input.ClipPreviews || j.input.ImageThumbnails; more; {
		if job.IsCancelled(ctx) {
			return
		}

		images, err := image.Query(ctx, r.Image, nil, findFilter)
		if err != nil {
			logger.Errorf("Error encountered queuing files to scan: %s", err.Error())
			return
		}

		for _, ss := range images {
			if job.IsCancelled(ctx) {
				return
			}

			if err := ss.LoadFiles(ctx, r.Image); err != nil {
				logger.Errorf("Error encountered queuing files to scan: %s", err.Error())
				return
			}

			j.queueImageJob(g, ss, queue)
		}

		if len(images) != batchSize {
			more = false
		} else {
			*findFilter.Page++
		}
	}
}

func getGeneratePreviewOptions(optionsInput GeneratePreviewOptionsInput) generate.PreviewOptions {
	config := config.GetInstance()

	ret := generate.PreviewOptions{
		Segments:        config.GetPreviewSegments(),
		SegmentDuration: config.GetPreviewSegmentDuration(),
		ExcludeStart:    config.GetPreviewExcludeStart(),
		ExcludeEnd:      config.GetPreviewExcludeEnd(),
		Preset:          config.GetPreviewPreset().String(),
		Audio:           config.GetPreviewAudio(),
	}

	if optionsInput.PreviewSegments != nil {
		ret.Segments = *optionsInput.PreviewSegments
	}

	if optionsInput.PreviewSegmentDuration != nil {
		ret.SegmentDuration = *optionsInput.PreviewSegmentDuration
	}

	if optionsInput.PreviewExcludeStart != nil {
		ret.ExcludeStart = *optionsInput.PreviewExcludeStart
	}

	if optionsInput.PreviewExcludeEnd != nil {
		ret.ExcludeEnd = *optionsInput.PreviewExcludeEnd
	}

	if optionsInput.PreviewPreset != nil {
		ret.Preset = optionsInput.PreviewPreset.String()
	}

	return ret
}

func (j *GenerateJob) queueSceneJobs(ctx context.Context, g *generate.Generator, scene *models.Scene, queue chan<- Task) {
	r := j.repository

	if j.input.Covers {
		task := &GenerateCoverTask{
			repository: r,
			Scene:      *scene,
			Overwrite:  j.overwrite,
		}

		if task.required(ctx) {
			j.totals.covers++
			j.totals.tasks++
			queue <- task
		}
	}

	if j.input.Sprites {
		task := &GenerateSpriteTask{
			Scene:               *scene,
			Overwrite:           j.overwrite,
			fileNamingAlgorithm: j.fileNamingAlgo,
		}

		if task.required() {
			j.totals.sprites++
			j.totals.tasks++
			queue <- task
		}
	}

	generatePreviewOptions := j.input.PreviewOptions
	if generatePreviewOptions == nil {
		generatePreviewOptions = &GeneratePreviewOptionsInput{}
	}
	options := getGeneratePreviewOptions(*generatePreviewOptions)

	if j.input.Previews {
		task := &GeneratePreviewTask{
			Scene:               *scene,
			ImagePreview:        j.input.ImagePreviews,
			Options:             options,
			Overwrite:           j.overwrite,
			fileNamingAlgorithm: j.fileNamingAlgo,
			generator:           g,
		}

		if task.required() {
			if task.videoPreviewRequired() {
				j.totals.previews++
			}
			if task.imagePreviewRequired() {
				j.totals.imagePreviews++
			}

			j.totals.tasks++
			queue <- task
		}
	}

	if j.input.Markers {
		task := &GenerateMarkersTask{
			repository:          r,
			Scene:               scene,
			Overwrite:           j.overwrite,
			fileNamingAlgorithm: j.fileNamingAlgo,
			ImagePreview:        j.input.MarkerImagePreviews,
			Screenshot:          j.input.MarkerScreenshots,

			generator: g,
		}

		markers := task.markersNeeded(ctx)
		if markers > 0 {
			j.totals.markers += int64(markers)
			j.totals.tasks++

			queue <- task
		}
	}

	if j.input.Transcodes {
		forceTranscode := j.input.ForceTranscodes
		task := &GenerateTranscodeTask{
			Scene:               *scene,
			Overwrite:           j.overwrite,
			Force:               forceTranscode,
			fileNamingAlgorithm: j.fileNamingAlgo,
			g:                   g,
		}
		if task.required() {
			j.totals.transcodes++
			j.totals.tasks++
			queue <- task
		}
	}

	if j.input.Phashes {
		// generate for all files in scene
		for _, f := range scene.Files.List() {
			task := &GeneratePhashTask{
				repository:          r,
				File:                f,
				fileNamingAlgorithm: j.fileNamingAlgo,
				Overwrite:           j.overwrite,
			}

			if task.required() {
				j.totals.phashes++
				j.totals.tasks++
				queue <- task
			}
		}
	}

	if j.input.InteractiveHeatmapsSpeeds {
		task := &GenerateInteractiveHeatmapSpeedTask{
			repository:          r,
			Scene:               *scene,
			Overwrite:           j.overwrite,
			fileNamingAlgorithm: j.fileNamingAlgo,
		}

		if task.required() {
			j.totals.interactiveHeatmapSpeeds++
			j.totals.tasks++
			queue <- task
		}
	}
}

func (j *GenerateJob) queueMarkerJob(g *generate.Generator, marker *models.SceneMarker, queue chan<- Task) {
	task := &GenerateMarkersTask{
		repository:          j.repository,
		Marker:              marker,
		Overwrite:           j.overwrite,
		fileNamingAlgorithm: j.fileNamingAlgo,
		generator:           g,
	}
	j.totals.markers++
	j.totals.tasks++
	queue <- task
}

func (j *GenerateJob) queueImageJob(g *generate.Generator, image *models.Image, queue chan<- Task) {
	if j.input.ImageThumbnails {
		task := &GenerateImageThumbnailTask{
			Image:     *image,
			Overwrite: j.overwrite,
		}

		if task.required() {
			j.totals.imageThumbnails++
			j.totals.tasks++
			queue <- task
		}
	}

	if j.input.ClipPreviews {
		task := &GenerateClipPreviewTask{
			Image:     *image,
			Overwrite: j.overwrite,
		}

		if task.required() {
			j.totals.clipPreviews++
			j.totals.tasks++
			queue <- task
		}
	}
}
