package main

import (
	"image"
	"image/color"
	"image/png"
	"math"
	"math/rand"
	"os"
	"path/filepath"
	"testing"
)

// --- synthetic scan image builder ---

type syntheticScan struct {
	img    *image.Gray
	dpi    float64
	cfg    FaceTemplateConfig
	layout pageLayout
}

func newSyntheticScan(dpi float64, cfg FaceTemplateConfig) *syntheticScan {
	if cfg.PageWidth <= 0 {
		cfg.PageWidth = 215.9
	}
	if cfg.PageHeight <= 0 {
		cfg.PageHeight = 279.4
	}
	if cfg.NumFaces < 1 {
		cfg.NumFaces = 6
	}
	if cfg.LongestSide <= 0 {
		cfg.LongestSide = 50
	}
	if cfg.GridSpacing <= 0 {
		cfg.GridSpacing = pickGridSpacing(cfg.LongestSide)
	}

	w := int(mmToPx(cfg.PageWidth, dpi))
	h := int(mmToPx(cfg.PageHeight, dpi))
	img := image.NewGray(image.Rect(0, 0, w, h))

	// White background
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			img.SetGray(x, y, color.Gray{Y: 240})
		}
	}

	layouts := computePageLayout(cfg)
	var layout pageLayout
	if len(layouts) > 0 {
		layout = layouts[0]
	}

	return &syntheticScan{img: img, dpi: dpi, cfg: cfg, layout: layout}
}

func (s *syntheticScan) drawRegistrationMarks() {
	margin := 15.0
	corners := [4][2]float64{
		{margin, margin},
		{s.cfg.PageWidth - margin, margin},
		{margin, s.cfg.PageHeight - margin},
		{s.cfg.PageWidth - margin, s.cfg.PageHeight - margin},
	}
	circR := mmToPx(0.8, s.dpi)
	for _, c := range corners {
		cx := mmToPx(c[0], s.dpi)
		cy := mmToPx(c[1], s.dpi)
		s.fillCircle(int(cx), int(cy), int(circR)+1, 0)
		// L-shaped arms
		armLen := int(mmToPx(5, s.dpi))
		s.drawHLine(int(cx), int(cy), armLen, 0)
		s.drawVLine(int(cx), int(cy), armLen, 0)
	}
}

func (s *syntheticScan) drawGridLinesInCell(cell faceCell) {
	grid := s.cfg.GridSpacing
	x0 := int(mmToPx(cell.X, s.dpi))
	y0 := int(mmToPx(cell.Y, s.dpi))
	x1 := int(mmToPx(cell.X+cell.W, s.dpi))
	y1 := int(mmToPx(cell.Y+cell.H, s.dpi))

	// Thin grid lines (very light, 1px)
	for gx := cell.X + grid; gx < cell.X+cell.W; gx += grid {
		px := int(mmToPx(gx, s.dpi))
		for py := y0; py < y1; py++ {
			s.setPixelSafe(px, py, 210)
		}
	}
	for gy := cell.Y + grid; gy < cell.Y+cell.H; gy += grid {
		py := int(mmToPx(gy, s.dpi))
		for px := x0; px < x1; px++ {
			s.setPixelSafe(px, py, 210)
		}
	}

	// Dashed cell border (medium darkness, thin)
	for px := x0; px < x1; px++ {
		if (px/6)%2 == 0 {
			s.setPixelSafe(px, y0, 80)
			s.setPixelSafe(px, y1, 80)
		}
	}
	for py := y0; py < y1; py++ {
		if (py/6)%2 == 0 {
			s.setPixelSafe(x0, py, 80)
			s.setPixelSafe(x1, py, 80)
		}
	}
}

// drawTracedPolygon draws a thick polygon in a cell (simulates pen tracing).
// Points are in mm relative to cell center.
func (s *syntheticScan) drawTracedPolygon(cell faceCell, outlineMM [][2]float64, thickness int) {
	cx := mmToPx(cell.X+cell.W/2, s.dpi)
	cy := mmToPx(cell.Y+cell.H/2, s.dpi)

	for i := 0; i < len(outlineMM); i++ {
		p0 := outlineMM[i]
		p1 := outlineMM[(i+1)%len(outlineMM)]
		x0 := cx + mmToPx(p0[0], s.dpi)
		y0 := cy + mmToPx(p0[1], s.dpi)
		x1 := cx + mmToPx(p1[0], s.dpi)
		y1 := cy + mmToPx(p1[1], s.dpi)
		s.drawThickLine(int(x0), int(y0), int(x1), int(y1), thickness, 20)
	}
}

func (s *syntheticScan) drawThickLine(x0, y0, x1, y1, thickness int, value uint8) {
	dx := float64(x1 - x0)
	dy := float64(y1 - y0)
	length := math.Sqrt(dx*dx + dy*dy)
	if length < 1 {
		return
	}
	steps := int(length * 2)
	for i := 0; i <= steps; i++ {
		t := float64(i) / float64(steps)
		px := float64(x0) + dx*t
		py := float64(y0) + dy*t
		for oy := -thickness; oy <= thickness; oy++ {
			for ox := -thickness; ox <= thickness; ox++ {
				if ox*ox+oy*oy <= thickness*thickness {
					s.setPixelSafe(int(px)+ox, int(py)+oy, value)
				}
			}
		}
	}
}

func (s *syntheticScan) fillCircle(cx, cy, r int, value uint8) {
	for dy := -r; dy <= r; dy++ {
		for dx := -r; dx <= r; dx++ {
			if dx*dx+dy*dy <= r*r {
				s.setPixelSafe(cx+dx, cy+dy, value)
			}
		}
	}
}

func (s *syntheticScan) drawHLine(x, y, length int, value uint8) {
	for dx := 0; dx < length; dx++ {
		s.setPixelSafe(x+dx, y, value)
	}
}

func (s *syntheticScan) drawVLine(x, y, length int, value uint8) {
	for dy := 0; dy < length; dy++ {
		s.setPixelSafe(x, y+dy, value)
	}
}

func (s *syntheticScan) setPixelSafe(x, y int, value uint8) {
	b := s.img.Bounds()
	if x >= b.Min.X && x < b.Max.X && y >= b.Min.Y && y < b.Max.Y {
		s.img.SetGray(x, y, color.Gray{Y: value})
	}
}

func (s *syntheticScan) addSaltPepperNoise(density float64, rng *rand.Rand) {
	b := s.img.Bounds()
	total := b.Dx() * b.Dy()
	nPixels := int(float64(total) * density)
	for i := 0; i < nPixels; i++ {
		x := rng.Intn(b.Dx()) + b.Min.X
		y := rng.Intn(b.Dy()) + b.Min.Y
		if rng.Intn(2) == 0 {
			s.img.SetGray(x, y, color.Gray{Y: 0})
		} else {
			s.img.SetGray(x, y, color.Gray{Y: 255})
		}
	}
}

func (s *syntheticScan) addExtraWhitespace(topPx, bottomPx, leftPx, rightPx int) *image.Gray {
	ob := s.img.Bounds()
	nw := ob.Dx() + leftPx + rightPx
	nh := ob.Dy() + topPx + bottomPx
	out := image.NewGray(image.Rect(0, 0, nw, nh))
	for y := 0; y < nh; y++ {
		for x := 0; x < nw; x++ {
			out.SetGray(x, y, color.Gray{Y: 255})
		}
	}
	for y := ob.Min.Y; y < ob.Max.Y; y++ {
		for x := ob.Min.X; x < ob.Max.X; x++ {
			out.SetGray(x+leftPx, y+topPx, s.img.GrayAt(x, y))
		}
	}
	return out
}

func (s *syntheticScan) rotate180() *image.Gray {
	b := s.img.Bounds()
	w, h := b.Dx(), b.Dy()
	out := image.NewGray(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			out.SetGray(w-1-x, h-1-y, s.img.GrayAt(x+b.Min.X, y+b.Min.Y))
		}
	}
	return out
}

func (s *syntheticScan) rotateSmall(angleDeg float64) *image.Gray {
	b := s.img.Bounds()
	w, h := b.Dx(), b.Dy()
	out := image.NewGray(image.Rect(0, 0, w, h))
	// Fill with white
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			out.SetGray(x, y, color.Gray{Y: 240})
		}
	}

	rad := angleDeg * math.Pi / 180
	cx, cy := float64(w)/2, float64(h)/2
	cosA, sinA := math.Cos(rad), math.Sin(rad)

	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			// inverse rotation to find source pixel
			dx := float64(x) - cx
			dy := float64(y) - cy
			srcX := int(dx*cosA+dy*sinA+cx+0.5) + b.Min.X
			srcY := int(-dx*sinA+dy*cosA+cy+0.5) + b.Min.Y
			if srcX >= b.Min.X && srcX < b.Max.X && srcY >= b.Min.Y && srcY < b.Max.Y {
				out.SetGray(x, y, s.img.GrayAt(srcX, srcY))
			}
		}
	}
	return out
}

func (s *syntheticScan) cropEdges(topMM, bottomMM, leftMM, rightMM float64) *image.Gray {
	b := s.img.Bounds()
	x0 := b.Min.X + int(mmToPx(leftMM, s.dpi))
	y0 := b.Min.Y + int(mmToPx(topMM, s.dpi))
	x1 := b.Max.X - int(mmToPx(rightMM, s.dpi))
	y1 := b.Max.Y - int(mmToPx(bottomMM, s.dpi))
	w, h := x1-x0, y1-y0
	out := image.NewGray(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			out.SetGray(x, y, s.img.GrayAt(x0+x, y0+y))
		}
	}
	return out
}

func saveTestImage(t *testing.T, img *image.Gray, dir, name string) string {
	t.Helper()
	path := filepath.Join(dir, name)
	f, err := os.Create(path)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	if err := png.Encode(f, img); err != nil {
		t.Fatal(err)
	}
	return path
}

// reference polygon: a diamond shape, 20mm across
var diamondMM = [][2]float64{
	{0, -10}, {10, 0}, {0, 10}, {-10, 0},
}

func makeCfg() FaceTemplateConfig {
	return FaceTemplateConfig{
		NumFaces: 4, LongestSide: 50,
		PageWidth: 215.9, PageHeight: 279.4,
	}
}

func buildCleanScan(t *testing.T, dpi float64) (*syntheticScan, string) {
	t.Helper()
	cfg := makeCfg()
	s := newSyntheticScan(dpi, cfg)
	s.drawRegistrationMarks()
	for _, cell := range s.layout.Cells {
		s.drawGridLinesInCell(cell)
	}
	// Draw diamond in face 1
	s.drawTracedPolygon(s.layout.Cells[0], diamondMM, int(mmToPx(1.0, dpi)))

	dir := t.TempDir()
	path := saveTestImage(t, s.img, dir, "scan.png")
	return s, path
}

func verifyDiamondExtracted(t *testing.T, faces []TracedFace) {
	t.Helper()
	if len(faces) == 0 {
		t.Fatal("no faces extracted")
	}
	// Find face 1
	var face1 *TracedFace
	for i := range faces {
		if faces[i].FaceNum == 1 {
			face1 = &faces[i]
			break
		}
	}
	if face1 == nil {
		t.Fatal("face 1 not found in extracted faces")
	}
	if len(face1.Outline) < 3 {
		t.Fatalf("face 1 has too few vertices: %d", len(face1.Outline))
	}

	// Check that the extracted polygon roughly fits a 20mm diamond
	var maxDim float64
	for _, pt := range face1.Outline {
		d := math.Sqrt(pt[0]*pt[0] + pt[1]*pt[1])
		if d > maxDim {
			maxDim = d
		}
	}
	// Diamond spans 10mm from center to tip; allow generous tolerance
	if maxDim < 5 || maxDim > 20 {
		t.Errorf("face 1 max dimension from center = %.1fmm, expected ~10mm", maxDim)
	}
}

// --- integration tests ---

func TestScanClean(t *testing.T) {
	cfg := makeCfg()
	_, path := buildCleanScan(t, 300)
	result, err := ProcessFaceScans([]string{path}, cfg)
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(result.DPI-300) > 5 {
		t.Errorf("DPI should be ~300, got %.0f", result.DPI)
	}
	verifyDiamondExtracted(t, result.Faces)
}

func TestScanWithNoise(t *testing.T) {
	cfg := makeCfg()
	s, _ := buildCleanScan(t, 300)
	rng := rand.New(rand.NewSource(42))
	s.addSaltPepperNoise(0.005, rng) // 0.5% noise

	dir := t.TempDir()
	path := saveTestImage(t, s.img, dir, "noisy.png")

	result, err := ProcessFaceScans([]string{path}, cfg)
	if err != nil {
		t.Fatal(err)
	}
	verifyDiamondExtracted(t, result.Faces)
}

func TestScanAskew(t *testing.T) {
	cfg := makeCfg()
	s, _ := buildCleanScan(t, 300)
	rotated := s.rotateSmall(1.5) // 1.5 degrees

	dir := t.TempDir()
	path := saveTestImage(t, rotated, dir, "askew.png")

	result, err := ProcessFaceScans([]string{path}, cfg)
	if err != nil {
		t.Fatal(err)
	}
	// Registration mark detection should compensate via affine transform
	verifyDiamondExtracted(t, result.Faces)
}

func TestScanExtraWhitespace(t *testing.T) {
	cfg := makeCfg()
	s, _ := buildCleanScan(t, 300)
	// Add 1 inch of whitespace on all sides
	padPx := int(mmToPx(25.4, 300))
	padded := s.addExtraWhitespace(padPx, padPx, padPx, padPx)

	dir := t.TempDir()
	path := saveTestImage(t, padded, dir, "padded.png")

	// DPI estimation will be off because image is larger than expected page,
	// but registration marks should still be found by searching near expected positions
	result, err := ProcessFaceScans([]string{path}, cfg)
	if err != nil {
		t.Fatal(err)
	}
	if len(result.Errors) == 0 {
		verifyDiamondExtracted(t, result.Faces)
	}
}

func TestScanCroppedEdges(t *testing.T) {
	cfg := makeCfg()
	s, _ := buildCleanScan(t, 300)
	// Crop 5mm off top and left — this removes two registration marks
	cropped := s.cropEdges(5, 0, 5, 0)

	dir := t.TempDir()
	path := saveTestImage(t, cropped, dir, "cropped.png")

	result, err := ProcessFaceScans([]string{path}, cfg)
	if err != nil {
		t.Fatal(err)
	}
	// Should still produce output (fallback to expected positions for missing marks)
	// May have errors logged but shouldn't crash
	if result == nil {
		t.Fatal("result should not be nil even with cropped edges")
	}
}

func TestScanUpsideDown(t *testing.T) {
	cfg := makeCfg()
	s, _ := buildCleanScan(t, 300)
	flipped := s.rotate180()

	dir := t.TempDir()
	path := saveTestImage(t, flipped, dir, "upsidedown.png")

	result, err := ProcessFaceScans([]string{path}, cfg)
	if err != nil {
		t.Fatal(err)
	}
	// The registration marks are symmetric (all 4 corners have identical marks)
	// so the affine transform should still find them, but the cell positions
	// will be mirrored. The system should either detect the flip or still extract
	// something (even if coordinates are mirrored).
	if result == nil {
		t.Fatal("result should not be nil for upside-down image")
	}
}

func TestScanHeavyNoise(t *testing.T) {
	cfg := makeCfg()
	s, _ := buildCleanScan(t, 300)
	rng := rand.New(rand.NewSource(99))
	s.addSaltPepperNoise(0.02, rng) // 2% noise — aggressive

	dir := t.TempDir()
	path := saveTestImage(t, s.img, dir, "heavy_noise.png")

	result, err := ProcessFaceScans([]string{path}, cfg)
	if err != nil {
		t.Fatal(err)
	}
	// May not extract perfect outlines but should not crash
	if result == nil {
		t.Fatal("result should not be nil with heavy noise")
	}
}

func TestScanLowDPI(t *testing.T) {
	cfg := makeCfg()
	s := newSyntheticScan(150, cfg) // low DPI
	s.drawRegistrationMarks()
	for _, cell := range s.layout.Cells {
		s.drawGridLinesInCell(cell)
	}
	s.drawTracedPolygon(s.layout.Cells[0], diamondMM, int(mmToPx(1.0, 150)))

	dir := t.TempDir()
	path := saveTestImage(t, s.img, dir, "lowdpi.png")

	result, err := ProcessFaceScans([]string{path}, cfg)
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(result.DPI-150) > 5 {
		t.Errorf("DPI should be ~150, got %.0f", result.DPI)
	}
	verifyDiamondExtracted(t, result.Faces)
}

func TestScanHighDPI(t *testing.T) {
	cfg := makeCfg()
	s := newSyntheticScan(600, cfg)
	s.drawRegistrationMarks()
	for _, cell := range s.layout.Cells {
		s.drawGridLinesInCell(cell)
	}
	s.drawTracedPolygon(s.layout.Cells[0], diamondMM, int(mmToPx(1.0, 600)))

	dir := t.TempDir()
	path := saveTestImage(t, s.img, dir, "highdpi.png")

	result, err := ProcessFaceScans([]string{path}, cfg)
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(result.DPI-600) > 10 {
		t.Errorf("DPI should be ~600, got %.0f", result.DPI)
	}
	verifyDiamondExtracted(t, result.Faces)
}

func TestScanMultiplePages(t *testing.T) {
	cfg := FaceTemplateConfig{
		NumFaces: 6, LongestSide: 120,
		PageWidth: 215.9, PageHeight: 279.4,
	}
	layouts := computePageLayout(cfg)
	if len(layouts) < 2 {
		t.Skip("config doesn't produce multiple pages")
	}

	var paths []string
	dir := t.TempDir()

	for pi, page := range layouts {
		s := newSyntheticScan(300, cfg)
		s.layout = page
		s.drawRegistrationMarks()
		for _, cell := range page.Cells {
			s.drawGridLinesInCell(cell)
			s.drawTracedPolygon(cell, diamondMM, int(mmToPx(1.0, 300)))
		}
		path := saveTestImage(t, s.img, dir, filepath.Base(t.Name())+string(rune('a'+pi))+".png")
		paths = append(paths, path)
	}

	result, err := ProcessFaceScans(paths, cfg)
	if err != nil {
		t.Fatal(err)
	}
	if len(result.Faces) < 2 {
		t.Errorf("expected faces from multiple pages, got %d", len(result.Faces))
	}
}