package main

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
)

const (
	trainingVideoMAEClipSeconds = 4.0
	trainingVideoMAENumFrames   = 16
	trainingVideoMAEFrameSize   = 224
	minVideoMAETrainCount       = 5
	minVideoMAEValCount         = 1
)

func trainingVideoMAEEpochs() int {
	raw := strings.TrimSpace(os.Getenv("TRAINING_VIDEOMAE_EPOCHS"))
	if raw == "" {
		return 8
	}

	n, err := strconv.Atoi(raw)
	if err != nil || n < 1 {
		return 8
	}
	if n > 200 {
		return 200
	}
	return n
}

func trainingVideoMAEBatchSize() int {
	raw := strings.TrimSpace(os.Getenv("TRAINING_VIDEOMAE_BATCH"))
	if raw == "" {
		return 2
	}

	n, err := strconv.Atoi(raw)
	if err != nil || n < 1 {
		return 2
	}
	if n > 32 {
		return 32
	}
	return n
}

func trainingVideoMAEEarlyStoppingPatience() int {
	epochs := trainingVideoMAEEpochs()
	patience := epochs / 3
	if patience < 2 {
		patience = 2
	}
	if patience > 8 {
		patience = 8
	}
	if patience > epochs {
		patience = epochs
	}
	return patience
}

type trainingVideoMAEManifestEntry struct {
	SampleID   string  `json:"sampleId"`
	Split      string  `json:"split"`
	Label      string  `json:"label"`
	ClipDir    string  `json:"clipDir"`
	SourcePath string  `json:"sourcePath,omitempty"`
	Second     float64 `json:"second,omitempty"`
	FrameCount int     `json:"frameCount"`
}

func trainingResolveVideoMAEModel(root string) trainingModelResolution {
	modelDir := filepath.Join(root, "videomae", "model")
	configPath := filepath.Join(modelDir, "config.json")
	if fileExistsNonEmpty(configPath) {
		return trainingModelResolution{
			BestPath:        modelDir,
			EffectivePath:   modelDir,
			Source:          "videomae_clip",
			TrainedExists:   true,
			EffectiveExists: true,
		}
	}

	return trainingModelResolution{
		BestPath:        modelDir,
		EffectivePath:   modelDir,
		Source:          "videomae_missing",
		TrainedExists:   false,
		EffectiveExists: false,
	}
}

func trainingEnsureVideoMAEDirs(root string) error {
	dirs := []string{
		filepath.Join(root, "videomae"),
		filepath.Join(root, "videomae", "dataset"),
		filepath.Join(root, "videomae", "dataset", "clips"),
		filepath.Join(root, "videomae", "dataset", "clips", "train"),
		filepath.Join(root, "videomae", "dataset", "clips", "val"),
		filepath.Join(root, "videomae", "model"),
		filepath.Join(root, "videomae", "runs"),
	}

	for _, dir := range dirs {
		if err := os.MkdirAll(dir, 0755); err != nil {
			return err
		}
	}

	return trainingWriteVideoMAELabelsFile(root)
}

func trainingVideoMAELabels() ([]string, error) {
	grouped, err := trainingGroupedLabels()
	if err != nil {
		return nil, err
	}

	out := []string{trainingNoSexPositionLabel}
	seen := map[string]bool{
		trainingNoSexPositionLabel: true,
	}

	for _, value := range grouped.SexPositions {
		label := normalizeSexPositionLabel(value)
		if seen[label] {
			continue
		}

		seen[label] = true
		out = append(out, label)
	}

	return out, nil
}

func trainingVideoMAELabelSet() (map[string]bool, error) {
	labels, err := trainingVideoMAELabels()
	if err != nil {
		return nil, err
	}

	out := map[string]bool{}
	for _, label := range labels {
		out[label] = true
	}
	return out, nil
}

func trainingWriteVideoMAELabelsFile(root string) error {
	labels, err := trainingVideoMAELabels()
	if err != nil {
		return err
	}

	body, err := json.MarshalIndent(map[string]any{
		"labels": labels,
	}, "", "  ")
	if err != nil {
		return err
	}

	return os.WriteFile(filepath.Join(root, "videomae", "dataset", "labels.json"), body, 0644)
}

func trainingVideoMAEManifestPath(root string) string {
	return filepath.Join(root, "videomae", "dataset", "manifest.jsonl")
}

func trainingVideoMAELabelForAnnotation(item TrainingAnnotation) string {
	if item.Negative {
		return trainingNoSexPositionLabel
	}

	effective := trainingEffectiveCorrection(item)
	label := normalizeSexPositionLabel(effective.SexPosition)
	if isNoSexPositionLabel(label) {
		return trainingNoSexPositionLabel
	}

	return label
}

func trainingVideoMAEFrameFallbackPath(root string, sampleID string) string {
	return filepath.Join(root, "frames", sampleID+".jpg")
}

func trainingVideoMAEAnnotationHasSource(root string, item TrainingAnnotation) bool {
	sourcePath := strings.TrimSpace(item.SourcePath)
	if sourcePath != "" && trainingSupportedImportVideo(sourcePath) && fileExistsNonEmpty(sourcePath) {
		return true
	}

	return fileExistsNonEmpty(trainingVideoMAEFrameFallbackPath(root, item.SampleID))
}

func trainingCountVideoMAEEligibleAnnotations(root string) (int, error) {
	items, err := trainingReadAnnotations(root)
	if err != nil {
		return 0, err
	}

	labelSet, err := trainingVideoMAELabelSet()
	if err != nil {
		return 0, err
	}

	count := 0
	for _, item := range items {
		sampleID := strings.TrimSpace(item.SampleID)
		if sampleID == "" || strings.Contains(sampleID, "/") || strings.Contains(sampleID, "\\") {
			continue
		}

		if !labelSet[trainingVideoMAELabelForAnnotation(item)] {
			continue
		}

		if trainingVideoMAEAnnotationHasSource(root, item) {
			count++
		}
	}

	return count, nil
}

func trainingReadVideoMAEManifest(root string) ([]trainingVideoMAEManifestEntry, error) {
	path := trainingVideoMAEManifestPath(root)
	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			return []trainingVideoMAEManifestEntry{}, nil
		}
		return nil, err
	}
	defer f.Close()

	out := []trainingVideoMAEManifestEntry{}
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}

		var entry trainingVideoMAEManifestEntry
		if err := json.Unmarshal([]byte(line), &entry); err != nil {
			continue
		}

		out = append(out, entry)
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return out, nil
}

func trainingWriteVideoMAEManifest(root string, entries []trainingVideoMAEManifestEntry) error {
	path := trainingVideoMAEManifestPath(root)
	if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		return err
	}

	var b strings.Builder
	for _, entry := range entries {
		line, err := json.Marshal(entry)
		if err != nil {
			return err
		}
		b.Write(line)
		b.WriteByte('\n')
	}

	return os.WriteFile(path, []byte(b.String()), 0644)
}

func trainingCountVideoMAEManifestSamples(root string) (train int, val int) {
	entries, err := trainingReadVideoMAEManifest(root)
	if err != nil {
		return 0, 0
	}

	for _, entry := range entries {
		if entry.FrameCount <= 0 || !fileExistsNonEmpty(filepath.Join(entry.ClipDir, "frame_001.jpg")) {
			continue
		}

		switch strings.ToLower(strings.TrimSpace(entry.Split)) {
		case "train":
			train++
		case "val":
			val++
		}
	}

	return train, val
}

func trainingRemoveVideoMAEGeneratedClips(root string) error {
	clipsDir := filepath.Join(root, "videomae", "dataset", "clips")
	if err := os.RemoveAll(clipsDir); err != nil {
		return err
	}

	for _, split := range []string{"train", "val"} {
		if err := os.MkdirAll(filepath.Join(clipsDir, split), 0755); err != nil {
			return err
		}
	}

	return nil
}

func trainingVideoMAEClipFrameCount(clipDir string) int {
	entries, err := os.ReadDir(clipDir)
	if err != nil {
		return 0
	}

	count := 0
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}

		ext := strings.ToLower(filepath.Ext(entry.Name()))
		if ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".webp" {
			count++
		}
	}

	return count
}

func trainingExtractVideoMAEClipFrames(
	ctx context.Context,
	videoPath string,
	centerSecond float64,
	clipDir string,
) error {
	if strings.TrimSpace(videoPath) == "" {
		return errors.New("video path missing")
	}

	if err := os.MkdirAll(clipDir, 0755); err != nil {
		return err
	}

	ffmpegPath := strings.TrimSpace(getSettings().FFmpegPath)
	if ffmpegPath == "" {
		ffmpegPath = "ffmpeg"
	}

	start := math.Max(0, centerSecond-trainingVideoMAEClipSeconds/2)
	fps := float64(trainingVideoMAENumFrames) / trainingVideoMAEClipSeconds
	vf := fmt.Sprintf(
		"fps=%.6f,scale=%d:%d:force_original_aspect_ratio=decrease,pad=%d:%d:(ow-iw)/2:(oh-ih)/2",
		fps,
		trainingVideoMAEFrameSize,
		trainingVideoMAEFrameSize,
		trainingVideoMAEFrameSize,
		trainingVideoMAEFrameSize,
	)

	cmd := exec.CommandContext(
		ctx,
		ffmpegPath,
		"-hide_banner",
		"-loglevel", "error",
		"-ss", fmt.Sprintf("%.3f", start),
		"-i", videoPath,
		"-t", fmt.Sprintf("%.3f", trainingVideoMAEClipSeconds),
		"-vf", vf,
		"-frames:v", strconv.Itoa(trainingVideoMAENumFrames),
		"-q:v", "3",
		filepath.Join(clipDir, "frame_%03d.jpg"),
	)

	hideCommandWindow(cmd)

	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("ffmpeg clip extract failed: %w: %s", err, strings.TrimSpace(string(out)))
	}

	if trainingVideoMAEClipFrameCount(clipDir) <= 0 {
		return errors.New("ffmpeg erzeugte keine Clip-Frames")
	}

	return nil
}

func trainingWriteStillVideoMAEClip(framePath string, clipDir string) error {
	if !fileExistsNonEmpty(framePath) {
		return errors.New("fallback frame missing")
	}

	if err := os.MkdirAll(clipDir, 0755); err != nil {
		return err
	}

	for i := 1; i <= trainingVideoMAENumFrames; i++ {
		dst := filepath.Join(clipDir, fmt.Sprintf("frame_%03d.jpg", i))
		if err := copyFile(framePath, dst); err != nil {
			return err
		}
	}

	return nil
}

func trainingWriteVideoMAEClipForAnnotation(
	ctx context.Context,
	root string,
	item TrainingAnnotation,
	clipDir string,
) (int, error) {
	if err := os.RemoveAll(clipDir); err != nil {
		return 0, err
	}
	if err := os.MkdirAll(clipDir, 0755); err != nil {
		return 0, err
	}

	sourcePath := strings.TrimSpace(item.SourcePath)
	if sourcePath != "" && trainingSupportedImportVideo(sourcePath) && fileExistsNonEmpty(sourcePath) {
		if err := trainingExtractVideoMAEClipFrames(ctx, sourcePath, item.Second, clipDir); err == nil {
			return trainingVideoMAEClipFrameCount(clipDir), nil
		} else {
			appLogln("videomae clip extract fallback:", item.SampleID, err)
		}
	}

	fallbackFrame := trainingVideoMAEFrameFallbackPath(root, item.SampleID)
	if err := trainingWriteStillVideoMAEClip(fallbackFrame, clipDir); err != nil {
		return 0, err
	}

	return trainingVideoMAEClipFrameCount(clipDir), nil
}

func trainingCopyVideoMAEClip(srcDir string, dstDir string) (int, error) {
	if err := os.RemoveAll(dstDir); err != nil {
		return 0, err
	}
	if err := os.MkdirAll(dstDir, 0755); err != nil {
		return 0, err
	}

	entries, err := os.ReadDir(srcDir)
	if err != nil {
		return 0, err
	}

	copied := 0
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}

		ext := strings.ToLower(filepath.Ext(entry.Name()))
		if ext != ".jpg" && ext != ".jpeg" && ext != ".png" && ext != ".webp" {
			continue
		}

		if err := copyFile(filepath.Join(srcDir, entry.Name()), filepath.Join(dstDir, entry.Name())); err != nil {
			return copied, err
		}
		copied++
	}

	return copied, nil
}

func trainingEnsureVideoMAEValidationEntries(
	root string,
	entries []trainingVideoMAEManifestEntry,
) []trainingVideoMAEManifestEntry {
	trainCount := 0
	valCount := 0
	trainEntries := []trainingVideoMAEManifestEntry{}

	for _, entry := range entries {
		switch strings.ToLower(strings.TrimSpace(entry.Split)) {
		case "train":
			trainCount++
			trainEntries = append(trainEntries, entry)
		case "val":
			valCount++
		}
	}

	if valCount >= minVideoMAEValCount || trainCount < minVideoMAETrainCount {
		return entries
	}

	sort.SliceStable(trainEntries, func(i, j int) bool {
		if trainEntries[i].Label == trainEntries[j].Label {
			return trainEntries[i].SampleID < trainEntries[j].SampleID
		}
		return trainEntries[i].Label < trainEntries[j].Label
	})

	for _, entry := range trainEntries {
		if valCount >= minVideoMAEValCount {
			break
		}

		copyID := entry.SampleID + "_valcopy"
		dstDir := filepath.Join(root, "videomae", "dataset", "clips", "val", copyID)
		frameCount, err := trainingCopyVideoMAEClip(entry.ClipDir, dstDir)
		if err != nil || frameCount <= 0 {
			if err != nil {
				appLogln("videomae val copy failed:", entry.SampleID, err)
			}
			continue
		}

		copyEntry := entry
		copyEntry.SampleID = copyID
		copyEntry.Split = "val"
		copyEntry.ClipDir = dstDir
		copyEntry.FrameCount = frameCount
		entries = append(entries, copyEntry)
		valCount++
	}

	return entries
}

func trainingSyncVideoMAEDataset(
	ctx context.Context,
	root string,
) (trainCount int, valCount int, written int, err error) {
	if err := trainingEnsureVideoMAEDirs(root); err != nil {
		return 0, 0, 0, err
	}
	if err := trainingRemoveVideoMAEGeneratedClips(root); err != nil {
		return 0, 0, 0, err
	}

	items, err := trainingReadAnnotations(root)
	if err != nil {
		return 0, 0, 0, err
	}

	labelSet, err := trainingVideoMAELabelSet()
	if err != nil {
		return 0, 0, 0, err
	}

	entries := []trainingVideoMAEManifestEntry{}

	for _, item := range items {
		if ctx.Err() != nil {
			return 0, 0, written, ctx.Err()
		}

		sampleID := strings.TrimSpace(item.SampleID)
		if sampleID == "" || strings.Contains(sampleID, "/") || strings.Contains(sampleID, "\\") {
			continue
		}

		label := trainingVideoMAELabelForAnnotation(item)
		if !labelSet[label] {
			continue
		}

		split := trainingStableSplit(sampleID)
		clipDir := filepath.Join(root, "videomae", "dataset", "clips", split, sampleID)
		frameCount, clipErr := trainingWriteVideoMAEClipForAnnotation(ctx, root, item, clipDir)
		if clipErr != nil {
			appLogln("videomae sample sync skipped:", sampleID, clipErr)
			continue
		}

		entry := trainingVideoMAEManifestEntry{
			SampleID:   sampleID,
			Split:      split,
			Label:      label,
			ClipDir:    clipDir,
			SourcePath: strings.TrimSpace(item.SourcePath),
			Second:     item.Second,
			FrameCount: frameCount,
		}
		entries = append(entries, entry)
		written++
	}

	entries = trainingEnsureVideoMAEValidationEntries(root, entries)

	for _, entry := range entries {
		switch strings.ToLower(strings.TrimSpace(entry.Split)) {
		case "train":
			trainCount++
		case "val":
			valCount++
		}
	}

	if err := trainingWriteVideoMAEManifest(root, entries); err != nil {
		return trainCount, valCount, written, err
	}

	return trainCount, valCount, written, nil
}