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 != "unknown" {
		t.Fatalf("sex position = %q, want unknown", effective.SexPosition)
	}
	if len(effective.Boxes) != 0 {
		t.Fatalf("negative annotation has %d boxes, want 0", len(effective.Boxes))
	}
}