package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestTrainingDetectorLabelContentAllowsExplicitNegative(t *testing.T) {
	content, err := trainingDetectorLabelContent(nil, map[string]int{"person": 0}, true)
	if err != nil {
		t.Fatalf("negative label content returned error: %v", err)
	}
	if len(content) != 0 {
		t.Fatalf("negative label content = %q, want empty", string(content))
	}
}

func TestTrainingDetectorLabelContentRejectsAccidentalEmptySample(t *testing.T) {
	_, err := trainingDetectorLabelContent(
		[]TrainingBox{{Label: "unknown_label", X: 0, Y: 0, W: 1, H: 1}},
		map[string]int{"person": 0},
		false,
	)
	if err == nil {
		t.Fatal("invalid non-negative sample should be rejected")
	}
}

func TestTrainingDetectorLabelContentWritesYOLOBox(t *testing.T) {
	content, err := trainingDetectorLabelContent(
		[]TrainingBox{{Label: "person", X: 0.1, Y: 0.2, W: 0.4, H: 0.6}},
		map[string]int{"person": 3},
		false,
	)
	if err != nil {
		t.Fatalf("positive label content returned error: %v", err)
	}

	got := strings.TrimSpace(string(content))
	want := "3 0.300000 0.500000 0.400000 0.600000"
	if got != want {
		t.Fatalf("label content = %q, want %q", got, want)
	}
}

func TestTrainingCountDetectorSamplesIncludesEmptyLabels(t *testing.T) {
	root := t.TempDir()
	imagesDir := filepath.Join(root, "images")
	labelsDir := filepath.Join(root, "labels")

	if err := os.MkdirAll(imagesDir, 0755); err != nil {
		t.Fatal(err)
	}
	if err := os.MkdirAll(labelsDir, 0755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(imagesDir, "negative.jpg"), []byte("image"), 0644); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(labelsDir, "negative.txt"), []byte{}, 0644); err != nil {
		t.Fatal(err)
	}

	if got := trainingCountDetectorSamples(imagesDir, labelsDir); got != 1 {
		t.Fatalf("sample count = %d, want 1", got)
	}
	if got := trainingCountPositiveDetectorSamples(imagesDir, labelsDir); got != 0 {
		t.Fatalf("positive sample count = %d, want 0", got)
	}
}

func TestTrainingEnsureDetectorValidationSampleIncludesPositiveExample(t *testing.T) {
	root := t.TempDir()
	trainImages := filepath.Join(root, "detector", "dataset", "images", "train")
	trainLabels := filepath.Join(root, "detector", "dataset", "labels", "train")

	if err := os.MkdirAll(trainImages, 0755); err != nil {
		t.Fatal(err)
	}
	if err := os.MkdirAll(trainLabels, 0755); err != nil {
		t.Fatal(err)
	}

	for i := 0; i < minDetectorTrainCount; i++ {
		id := fmt.Sprintf("sample-%02d", i)
		if err := os.WriteFile(filepath.Join(trainImages, id+".jpg"), []byte("image"), 0644); err != nil {
			t.Fatal(err)
		}

		label := []byte{}
		if i == minDetectorTrainCount-1 {
			label = []byte("0 0.5 0.5 1 1\n")
		}
		if err := os.WriteFile(filepath.Join(trainLabels, id+".txt"), label, 0644); err != nil {
			t.Fatal(err)
		}
	}

	if err := trainingEnsureDetectorValidationSample(root); err != nil {
		t.Fatal(err)
	}

	valImages := filepath.Join(root, "detector", "dataset", "images", "val")
	valLabels := filepath.Join(root, "detector", "dataset", "labels", "val")
	if got := trainingCountDetectorSamples(valImages, valLabels); got < minDetectorValCount {
		t.Fatalf("validation sample count = %d, want at least %d", got, minDetectorValCount)
	}
	if got := trainingCountPositiveDetectorSamples(valImages, valLabels); got < 1 {
		t.Fatalf("positive validation sample count = %d, want at least 1", got)
	}
}

func TestTrainingEffectiveCorrectionClearsNegativeAnnotation(t *testing.T) {
	effective := trainingEffectiveCorrection(TrainingAnnotation{
		Negative: true,
		Prediction: TrainingPrediction{
			SexPosition: "doggy",
			Boxes: []TrainingBox{
				{Label: "person", X: 0, Y: 0, W: 1, H: 1},
			},
		},
	})

	if effective.SexPosition != trainingNoSexPositionLabel {
		t.Fatalf("sex position = %q, want %s", effective.SexPosition, trainingNoSexPositionLabel)
	}
	if len(effective.Boxes) != 0 {
		t.Fatalf("negative annotation has %d boxes, want 0", len(effective.Boxes))
	}
}

func trainingTestPoseKeypoints(overrides map[string]TrainingKeypoint) []TrainingKeypoint {
	names := []string{
		"nose",
		"left_eye", "right_eye",
		"left_ear", "right_ear",
		"left_shoulder", "right_shoulder",
		"left_elbow", "right_elbow",
		"left_wrist", "right_wrist",
		"left_hip", "right_hip",
		"left_knee", "right_knee",
		"left_ankle", "right_ankle",
	}

	out := make([]TrainingKeypoint, 0, len(names))
	for i, name := range names {
		point := TrainingKeypoint{
			Name: name,
			X:    0.20 + float64(i)*0.01,
			Y:    0.25 + float64(i)*0.01,
			Conf: 0.85,
		}

		if override, ok := overrides[name]; ok {
			override.Name = name
			point = override
		}

		out = append(out, point)
	}

	return out
}

func TestTrainingApplyPoseToPredictionAggregatesMultiplePersons(t *testing.T) {
	pred := TrainingPrediction{
		SexPosition: trainingNoSexPositionLabel,
		Boxes: []TrainingBox{
			{Label: "person_female", X: 0.10, Y: 0.10, W: 0.35, H: 0.75},
			{Label: "person_male", X: 0.42, Y: 0.12, W: 0.34, H: 0.72},
		},
	}

	pose := TrainingPosePrediction{
		Available: true,
		Persons: []TrainingPosePerson{
			{
				Label:     "missionary",
				Score:     0.58,
				Box:       TrainingBox{X: 0.10, Y: 0.10, W: 0.35, H: 0.75},
				Keypoints: trainingTestPoseKeypoints(nil),
			},
			{
				Label:     "doggy",
				Score:     0.55,
				Box:       TrainingBox{X: 0.42, Y: 0.12, W: 0.34, H: 0.72},
				Keypoints: trainingTestPoseKeypoints(nil),
			},
			{
				Label:     "doggy",
				Score:     0.54,
				Box:       TrainingBox{X: 0.52, Y: 0.14, W: 0.30, H: 0.70},
				Keypoints: trainingTestPoseKeypoints(nil),
			},
		},
	}

	got := trainingApplyPoseToPrediction(pred, pose)
	if got.SexPosition != "doggy" {
		t.Fatalf("sex position = %q, want doggy", got.SexPosition)
	}
}

func TestTrainingApplyPoseToPredictionUsesObjectPositionContext(t *testing.T) {
	pred := TrainingPrediction{
		SexPosition: trainingNoSexPositionLabel,
		Boxes: []TrainingBox{
			{Label: "person_female", X: 0.08, Y: 0.08, W: 0.42, H: 0.78},
			{Label: "person_male", X: 0.42, Y: 0.12, W: 0.35, H: 0.76},
			{Label: "penis", X: 0.50, Y: 0.48, W: 0.07, H: 0.08},
		},
	}

	pose := TrainingPosePrediction{
		Available: true,
		Persons: []TrainingPosePerson{
			{
				Label: "missionary",
				Score: 0.58,
				Box:   TrainingBox{X: 0.08, Y: 0.08, W: 0.42, H: 0.78},
				Keypoints: trainingTestPoseKeypoints(map[string]TrainingKeypoint{
					"nose": {X: 0.18, Y: 0.18, Conf: 0.90},
				}),
			},
			{
				Label: "blowjob",
				Score: 0.55,
				Box:   TrainingBox{X: 0.42, Y: 0.12, W: 0.35, H: 0.76},
				Keypoints: trainingTestPoseKeypoints(map[string]TrainingKeypoint{
					"nose": {X: 0.53, Y: 0.52, Conf: 0.95},
				}),
			},
		},
	}

	got := trainingApplyPoseToPrediction(pred, pose)
	if got.SexPosition != "blowjob" {
		t.Fatalf("sex position = %q, want blowjob", got.SexPosition)
	}
}

func TestTrainingApplyPoseToPredictionUsesOcclusionTolerantCowgirlGeometry(t *testing.T) {
	pred := TrainingPrediction{
		SexPosition: trainingNoSexPositionLabel,
		Boxes: []TrainingBox{
			{Label: "person_female", X: 0.30, Y: 0.10, W: 0.28, H: 0.62},
			{Label: "person_male", X: 0.20, Y: 0.42, W: 0.55, H: 0.24},
		},
	}

	pose := TrainingPosePrediction{
		Available: true,
		Persons: []TrainingPosePerson{
			{
				Label: "person",
				Score: 0.72,
				Box:   TrainingBox{X: 0.30, Y: 0.10, W: 0.28, H: 0.62},
				Keypoints: trainingTestPoseKeypoints(map[string]TrainingKeypoint{
					"left_shoulder":  {X: 0.40, Y: 0.22, Conf: 0.92},
					"right_shoulder": {X: 0.48, Y: 0.22, Conf: 0.92},
					"left_hip":       {X: 0.38, Y: 0.38, Conf: 0.92},
					"right_hip":      {X: 0.50, Y: 0.38, Conf: 0.92},
					"left_knee":      {X: 0.31, Y: 0.58, Conf: 0.90},
					"right_knee":     {X: 0.57, Y: 0.58, Conf: 0.90},
				}),
			},
			{
				Label: "person",
				Score: 0.70,
				Box:   TrainingBox{X: 0.20, Y: 0.42, W: 0.55, H: 0.24},
				Keypoints: trainingTestPoseKeypoints(map[string]TrainingKeypoint{
					"left_shoulder":  {X: 0.25, Y: 0.50, Conf: 0.90},
					"right_shoulder": {X: 0.38, Y: 0.50, Conf: 0.90},
					"left_hip":       {X: 0.52, Y: 0.53, Conf: 0.90},
					"right_hip":      {X: 0.66, Y: 0.53, Conf: 0.90},
					"left_knee":      {X: 0.58, Y: 0.56, Conf: 0.88},
					"right_knee":     {X: 0.70, Y: 0.56, Conf: 0.88},
				}),
			},
		},
	}

	got := trainingApplyPoseToPrediction(pred, pose)
	if got.SexPosition != "cowgirl" {
		t.Fatalf("sex position = %q, want cowgirl", got.SexPosition)
	}
	if got.SexPositionScore < 0.30 {
		t.Fatalf("sex position score = %.3f, want >= 0.30", got.SexPositionScore)
	}
}

func TestTrainingApplyPoseToPredictionKeepsUnreliablePoseOutOfContext(t *testing.T) {
	pred := TrainingPrediction{
		SexPosition: trainingNoSexPositionLabel,
		Boxes: []TrainingBox{
			{Label: "person_female", X: 0.08, Y: 0.08, W: 0.42, H: 0.78},
			{Label: "penis", X: 0.50, Y: 0.48, W: 0.07, H: 0.08},
		},
	}

	pose := TrainingPosePrediction{
		Available: true,
		Persons: []TrainingPosePerson{
			{
				Label: "blowjob",
				Score: 0.24,
				Box:   TrainingBox{X: 0.08, Y: 0.08, W: 0.42, H: 0.78},
				Keypoints: trainingTestPoseKeypoints(map[string]TrainingKeypoint{
					"nose": {X: 0.53, Y: 0.52, Conf: 0.95},
				}),
			},
		},
	}

	got := trainingApplyPoseToPrediction(pred, pose)
	if got.SexPosition != trainingNoSexPositionLabel {
		t.Fatalf("sex position = %q, want %s", got.SexPosition, trainingNoSexPositionLabel)
	}
	if len(got.Persons) != 1 {
		t.Fatalf("persons = %d, want 1", len(got.Persons))
	}
	if got.Persons[0].Reliable {
		t.Fatalf("low-score pose should be marked unreliable: %+v", got.Persons[0])
	}
	if got.Persons[0].VisibleKeypoints == 0 || got.Persons[0].Quality == 0 {
		t.Fatalf("pose quality should be annotated: %+v", got.Persons[0])
	}
}

func TestTrainingApplyPoseToPredictionUsesBoxContextWithoutPose(t *testing.T) {
	pred := TrainingPrediction{
		ModelAvailable: true,
		Source:         "yolo26_detector",
		SexPosition:    trainingNoSexPositionLabel,
		Boxes: []TrainingBox{
			{Label: "person_female", X: 0.18, Y: 0.12, W: 0.42, H: 0.72},
			{Label: "person_male", X: 0.42, Y: 0.18, W: 0.38, H: 0.68},
			{Label: "penis", X: 0.48, Y: 0.50, W: 0.08, H: 0.08},
			{Label: "pussy", X: 0.54, Y: 0.51, W: 0.08, H: 0.08},
		},
	}

	got := trainingApplyPoseToPrediction(pred, TrainingPosePrediction{Available: false})
	if got.SexPosition == trainingNoSexPositionLabel {
		t.Fatalf("sex position = %q, want uncertain context prediction", got.SexPosition)
	}
	if got.SexPositionScore < trainingPositionContextMinScore {
		t.Fatalf("sex position score = %.3f, want >= %.3f", got.SexPositionScore, trainingPositionContextMinScore)
	}
	if got.SexPositionScore > trainingPositionContextMaxScore {
		t.Fatalf("sex position score = %.3f, want <= %.3f", got.SexPositionScore, trainingPositionContextMaxScore)
	}
	if !strings.Contains(got.Source, "box_context") {
		t.Fatalf("source = %q, want box_context", got.Source)
	}
}

func TestTrainingApplyPoseToPredictionBoostsPoseWithBoxContext(t *testing.T) {
	pred := TrainingPrediction{
		ModelAvailable: true,
		Source:         "yolo26_detector",
		SexPosition:    trainingNoSexPositionLabel,
		Boxes: []TrainingBox{
			{Label: "penis", X: 0.48, Y: 0.50, W: 0.08, H: 0.08},
			{Label: "pussy", X: 0.54, Y: 0.51, W: 0.08, H: 0.08},
		},
	}

	pose := TrainingPosePrediction{
		Available: true,
		Persons: []TrainingPosePerson{
			{
				Label:     "doggy",
				Score:     0.31,
				Box:       TrainingBox{X: 0.20, Y: 0.12, W: 0.60, H: 0.76},
				Keypoints: trainingTestPoseKeypoints(nil),
			},
		},
	}

	got := trainingApplyPoseToPrediction(pred, pose)
	if got.SexPosition != "doggy" {
		t.Fatalf("sex position = %q, want doggy", got.SexPosition)
	}
	if got.SexPositionScore <= 0.31 {
		t.Fatalf("sex position score = %.3f, want context boost over 0.31", got.SexPositionScore)
	}
	if !strings.Contains(got.Source, "yolo_pose") || !strings.Contains(got.Source, "box_context") {
		t.Fatalf("source = %q, want yolo_pose and box_context", got.Source)
	}
}

func TestBuildClipPositionHitsFromEvidenceRequiresTemporalSupport(t *testing.T) {
	hits := buildClipPositionHitsFromEvidence([]analyzePositionEvidence{
		{Time: 1, Label: "doggy", Score: 0.55, HasPose: true, PersonCount: 2},
		{Time: 2, Label: "cowgirl", Score: 0.60, HasPose: true, PersonCount: 2},
		{Time: 6, Label: "doggy", Score: 0.58, HasPose: true, PersonCount: 2},
	}, 10)

	if len(hits) != 0 {
		t.Fatalf("hits = %+v, want no temporally supported position", hits)
	}
}

func TestBuildClipPositionHitsFromEvidenceBuildsStablePosition(t *testing.T) {
	hits := buildClipPositionHitsFromEvidence([]analyzePositionEvidence{
		{Time: 1, Label: "doggy", Score: 0.45, HasPose: true, HasContext: true, PersonCount: 2},
		{Time: 2, Label: "doggy", Score: 0.47, HasPose: true, HasContext: true, PersonCount: 2},
		{Time: 3, Label: "doggy", Score: 0.44, HasPose: true, HasContext: true, PersonCount: 2},
		{Time: 6, Label: "cowgirl", Score: 0.62, HasPose: true, PersonCount: 2},
	}, 10)

	if len(hits) == 0 {
		t.Fatal("expected stable doggy position hit")
	}
	if hits[0].Label != "position:doggy" {
		t.Fatalf("label = %q, want position:doggy", hits[0].Label)
	}
	if hits[0].End <= hits[0].Start {
		t.Fatalf("invalid hit span: %+v", hits[0])
	}
}

func TestBuildClipPositionHitsFromEvidenceCapsContextOnlyScore(t *testing.T) {
	hits := buildClipPositionHitsFromEvidence([]analyzePositionEvidence{
		{Time: 1, Label: "missionary", Score: 0.90, HasContext: true, PersonCount: 2},
		{Time: 2, Label: "missionary", Score: 0.88, HasContext: true, PersonCount: 2},
		{Time: 3, Label: "missionary", Score: 0.92, HasContext: true, PersonCount: 2},
	}, 8)

	if len(hits) == 0 {
		t.Fatal("expected context-only position hit")
	}
	if hits[0].Score > trainingPositionContextMaxScore {
		t.Fatalf("score = %.3f, want <= %.3f", hits[0].Score, trainingPositionContextMaxScore)
	}
}

func TestTrainingFilterPosePersonsByContextDropsUnannotatedPeople(t *testing.T) {
	persons := []TrainingPosePerson{
		{
			Label:     "doggy",
			Score:     0.80,
			Box:       TrainingBox{X: 0.10, Y: 0.10, W: 0.30, H: 0.70},
			Keypoints: trainingTestPoseKeypoints(nil),
		},
		{
			Label:     "doggy",
			Score:     0.70,
			Box:       TrainingBox{X: 0.70, Y: 0.10, W: 0.20, H: 0.70},
			Keypoints: trainingTestPoseKeypoints(nil),
		},
	}

	filtered := trainingFilterPosePersonsByContext(persons, []TrainingBox{
		{Label: "person_female", X: 0.08, Y: 0.08, W: 0.34, H: 0.74},
	})

	if len(filtered) != 1 {
		t.Fatalf("filtered persons = %d, want 1", len(filtered))
	}
	if filtered[0].Box.X != persons[0].Box.X {
		t.Fatalf("kept wrong person: %+v", filtered[0].Box)
	}
}