package imageproc

import (
	"image"
	"math"
	"sort"

	"materialanalyzer/internal/model"
)

type rgbSample struct {
	r, g, b float64
	lum     float64
	sat     float64
}

type colorCluster struct {
	center rgbSample
	count  int
}

func ExtractVisual(img image.Image) (model.VisualFeatures, model.ColorHistogram) {
	samples := collectSamples(img, 65000)
	hist := buildHSVHistogram(samples, 16, 4, 4)
	if len(samples) == 0 {
		return model.VisualFeatures{HistogramBins: 256, HistogramSpace: "hsv"}, hist
	}

	var sumR, sumG, sumB, sumLum, sumLumSq, sumSat float64
	for _, s := range samples {
		sumR += s.r
		sumG += s.g
		sumB += s.b
		sumLum += s.lum
		sumLumSq += s.lum * s.lum
		sumSat += s.sat
	}
	n := float64(len(samples))
	avg := model.RGB{R: int(math.Round(sumR / n)), G: int(math.Round(sumG / n)), B: int(math.Round(sumB / n))}
	avgHSV := RGBToHSV(avg)
	avgLAB := RGBToLAB(avg)
	brightness := sumLum / n / 255.0
	contrast := math.Sqrt(math.Max(0, sumLumSq/n-(sumLum/n)*(sumLum/n))) / 255.0
	saturation := sumSat / n

	clusters := kmeansRGB(samples, 5, 12)
	colorClusters := make([]model.ColorCluster, 0, len(clusters))
	for _, c := range clusters {
		if c.count == 0 {
			continue
		}
		rgb := model.RGB{
			R: int(math.Round(c.center.r)),
			G: int(math.Round(c.center.g)),
			B: int(math.Round(c.center.b)),
		}
		colorClusters = append(colorClusters, model.ColorCluster{
			RGB: rgb, LAB: RGBToLAB(rgb), HSV: RGBToHSV(rgb),
			Percentage: round4(float64(c.count) / n),
		})
	}
	sort.Slice(colorClusters, func(i, j int) bool { return colorClusters[i].Percentage > colorClusters[j].Percentage })

	dominant := avg
	secondary := avg
	if len(colorClusters) > 0 {
		dominant = colorClusters[0].RGB
	}
	if len(colorClusters) > 1 {
		secondary = colorClusters[1].RGB
	}

	return model.VisualFeatures{
		DominantRGB: dominant, SecondaryRGB: secondary,
		DominantLAB: RGBToLAB(dominant), DominantHSV: RGBToHSV(dominant),
		AverageRGB: avg, AverageLAB: avgLAB, AverageHSV: avgHSV,
		Brightness: round4(brightness), Contrast: round4(contrast), Saturation: round4(saturation),
		ColorClusters: colorClusters,
		HistogramBins: 256, HistogramSpace: "hsv",
	}, hist
}

func collectSamples(img image.Image, maxSamples int) []rgbSample {
	b := img.Bounds()
	total := b.Dx() * b.Dy()
	if total <= 0 {
		return nil
	}
	step := 1
	if total > maxSamples {
		step = int(math.Ceil(math.Sqrt(float64(total) / float64(maxSamples))))
	}
	samples := make([]rgbSample, 0, min(total/(step*step)+1, maxSamples+1024))
	for y := b.Min.Y; y < b.Max.Y; y += step {
		for x := b.Min.X; x < b.Max.X; x += step {
			r, g, bb := rgba8(img.At(x, y))
			hsv := rgbToHSVFloat(float64(r), float64(g), float64(bb))
			lum := 0.2126*float64(r) + 0.7152*float64(g) + 0.0722*float64(bb)
			samples = append(samples, rgbSample{r: float64(r), g: float64(g), b: float64(bb), lum: lum, sat: hsv.S})
		}
	}
	return samples
}

func buildHSVHistogram(samples []rgbSample, hueBins, saturationBins, valueBins int) model.ColorHistogram {
	values := make([]float64, hueBins*saturationBins*valueBins)
	if len(samples) == 0 {
		return model.ColorHistogram{Space: "hsv", Bins: []int{hueBins, saturationBins, valueBins}, Values: values}
	}
	for _, s := range samples {
		hsv := rgbToHSVFloat(s.r, s.g, s.b)
		hi := min(hueBins-1, int(hsv.H/360.0*float64(hueBins)))
		si := min(saturationBins-1, int(hsv.S*float64(saturationBins)))
		vi := min(valueBins-1, int(hsv.V*float64(valueBins)))
		values[hi*saturationBins*valueBins+si*valueBins+vi]++
	}
	denom := float64(len(samples))
	for i := range values {
		values[i] = round6(values[i] / denom)
	}
	return model.ColorHistogram{Space: "hsv", Bins: []int{hueBins, saturationBins, valueBins}, Values: values}
}

func kmeansRGB(samples []rgbSample, k, iterations int) []colorCluster {
	if len(samples) == 0 || k <= 0 {
		return nil
	}
	sorted := append([]rgbSample(nil), samples...)
	sort.Slice(sorted, func(i, j int) bool { return sorted[i].lum < sorted[j].lum })
	if k > len(sorted) {
		k = len(sorted)
	}
	centers := make([]rgbSample, k)
	for i := 0; i < k; i++ {
		idx := int((float64(i) + 0.5) / float64(k) * float64(len(sorted)-1))
		centers[i] = sorted[idx]
	}

	assignments := make([]int, len(samples))
	for iter := 0; iter < iterations; iter++ {
		sums := make([]rgbSample, k)
		counts := make([]int, k)
		for i, s := range samples {
			bestIdx := 0
			bestDist := math.MaxFloat64
			for c, center := range centers {
				dr, dg, db := s.r-center.r, s.g-center.g, s.b-center.b
				dist := dr*dr + dg*dg + db*db
				if dist < bestDist {
					bestDist = dist
					bestIdx = c
				}
			}
			assignments[i] = bestIdx
			sums[bestIdx].r += s.r
			sums[bestIdx].g += s.g
			sums[bestIdx].b += s.b
			sums[bestIdx].lum += s.lum
			sums[bestIdx].sat += s.sat
			counts[bestIdx]++
		}
		for i := range centers {
			if counts[i] == 0 {
				continue
			}
			denom := float64(counts[i])
			centers[i] = rgbSample{
				r: sums[i].r / denom, g: sums[i].g / denom, b: sums[i].b / denom,
				lum: sums[i].lum / denom, sat: sums[i].sat / denom,
			}
		}
	}

	clusters := make([]colorCluster, k)
	for i := range clusters {
		clusters[i].center = centers[i]
	}
	for _, idx := range assignments {
		clusters[idx].count++
	}
	sort.Slice(clusters, func(i, j int) bool { return clusters[i].count > clusters[j].count })
	return clusters
}

func RGBToHSV(rgb model.RGB) model.HSV {
	return rgbToHSVFloat(float64(rgb.R), float64(rgb.G), float64(rgb.B))
}

func rgbToHSVFloat(r, g, b float64) model.HSV {
	r /= 255
	g /= 255
	b /= 255
	maxV := math.Max(r, math.Max(g, b))
	minV := math.Min(r, math.Min(g, b))
	d := maxV - minV
	h := 0.0
	if d != 0 {
		switch maxV {
		case r:
			h = math.Mod((g-b)/d, 6)
		case g:
			h = (b-r)/d + 2
		default:
			h = (r-g)/d + 4
		}
		h *= 60
		if h < 0 {
			h += 360
		}
	}
	s := 0.0
	if maxV != 0 {
		s = d / maxV
	}
	return model.HSV{H: round4(h), S: round4(s), V: round4(maxV)}
}

func RGBToLAB(rgb model.RGB) model.LAB {
	r := pivotRGB(float64(rgb.R) / 255.0)
	g := pivotRGB(float64(rgb.G) / 255.0)
	b := pivotRGB(float64(rgb.B) / 255.0)

	x := (r*0.4124 + g*0.3576 + b*0.1805) / 0.95047
	y := (r*0.2126 + g*0.7152 + b*0.0722) / 1.00000
	z := (r*0.0193 + g*0.1192 + b*0.9505) / 1.08883

	fx, fy, fz := pivotXYZ(x), pivotXYZ(y), pivotXYZ(z)
	return model.LAB{
		L: round4(116*fy - 16),
		A: round4(500 * (fx - fy)),
		B: round4(200 * (fy - fz)),
	}
}

func pivotRGB(v float64) float64 {
	if v > 0.04045 {
		return math.Pow((v+0.055)/1.055, 2.4)
	}
	return v / 12.92
}

func pivotXYZ(v float64) float64 {
	if v > 0.008856 {
		return math.Cbrt(v)
	}
	return 7.787*v + 16.0/116.0
}

func round4(v float64) float64 { return math.Round(v*10000) / 10000 }
func round6(v float64) float64 { return math.Round(v*1000000) / 1000000 }