package photos

import (
	"image"
	"math"
	"runtime"
	"sync"

	"github.com/disintegration/imaging"
)

// Spatial effects of the non-destructive editor. They run after the light and
// color adjustments, in a fixed order: dehaze, texture, clarity, vignette,
// grain. Texture and clarity compute a correction from the luminance and add it
// to all three channels equally, so hue and saturation are preserved; dehaze is
// a global tonal correction based on the dark channel prior. Vignette and grain
// are the final, "post-crop" operations.
//
// Everything here is deterministic: the grain seed comes from the photo id, so
// a repeated render of the same parameters produces the same pixels.

// maxDehazeSide is the longest side of the reduced copy used to estimate the
// atmosphere and the transmission map.
const maxDehazeSide = 512

// dehazeOmega is the share of the haze the dark channel prior removes.
const dehazeOmega = 0.95

// dehazeMinT bounds the transmission map from below, so the darkest regions do
// not amplify noise without limit.
const dehazeMinT = 0.1

// applyEffects applies the non-empty effect passes in the documented order and
// returns the image owned by the caller.
func applyEffects(dst *image.NRGBA, p EffectsParams, seed uint64) *image.NRGBA {
	if p.Dehaze != 0 {
		dst = applyDehaze(dst, p.Dehaze)
	}
	if p.Texture != 0 {
		applyTexture(dst, p.Texture)
	}
	if p.Clarity != 0 {
		applyClarity(dst, p.Clarity)
	}
	if p.Vignette.Amount != 0 {
		applyVignette(dst, p.Vignette)
	}
	if p.Grain.Amount != 0 {
		applyGrain(dst, p.Grain, seed)
	}
	return dst
}

// luminance returns the perceived brightness Y of an RGB pixel (each 0..255).
func luminance(r, g, b float64) float64 {
	return 0.299*r + 0.587*g + 0.114*b
}

// lumaPlane builds a Y plane of the image, one byte per pixel.
func lumaPlane(img *image.NRGBA) []uint8 {
	w, h := img.Rect.Dx(), img.Rect.Dy()
	out := make([]uint8, w*h)
	parallelChunks(h, func(from, to int) {
		for y := from; y < to; y++ {
			row := img.Pix[y*img.Stride : y*img.Stride+w*4]
			dst := out[y*w : (y+1)*w]
			for x := 0; x < w; x++ {
				i := x * 4
				dst[x] = uint8(luminance(float64(row[i]), float64(row[i+1]), float64(row[i+2])) + 0.5)
			}
		}
	})
	return out
}

// parallelChunks splits n independent items (rows or columns) into chunks and
// runs fn over each chunk, mirroring the row parallelism of applyLight and
// applyColor.
func parallelChunks(n int, fn func(from, to int)) {
	if n <= 0 {
		return
	}
	workers := runtime.GOMAXPROCS(0)
	if workers > n {
		workers = n
	}
	if workers <= 1 {
		fn(0, n)
		return
	}
	chunk := (n + workers - 1) / workers
	var wg sync.WaitGroup
	for from := 0; from < n; from += chunk {
		to := from + chunk
		if to > n {
			to = n
		}
		wg.Add(1)
		go func(from, to int) {
			defer wg.Done()
			fn(from, to)
		}(from, to)
	}
	wg.Wait()
}

// boxBlur applies one separable box blur with edge-clamped borders. The running
// window makes the cost independent of the radius. Edge pixels repeat the
// nearest border value, so the size of the image never changes.
func boxBlur(src []uint8, w, h, radius int) []uint8 {
	if w <= 0 || h <= 0 || len(src) < w*h {
		return nil
	}
	out := make([]uint8, w*h)
	if radius <= 0 {
		copy(out, src[:w*h])
		return out
	}
	tmp := make([]uint8, w*h)
	blurHorizontal(src, tmp, w, h, radius)
	blurVertical(tmp, out, w, h, radius)
	return out
}

// boxBlur3 approximates a Gaussian blur with three box passes.
func boxBlur3(src []uint8, w, h, radius int) []uint8 {
	out := boxBlur(src, w, h, radius)
	out = boxBlur(out, w, h, radius)
	return boxBlur(out, w, h, radius)
}

// blurHorizontal is one horizontal box-blur pass with a sliding window.
func blurHorizontal(src, dst []uint8, w, h, radius int) {
	window := 2*radius + 1
	parallelChunks(h, func(from, to int) {
		for y := from; y < to; y++ {
			row := src[y*w : (y+1)*w]
			out := dst[y*w : (y+1)*w]
			var sum int
			for k := -radius; k <= radius; k++ {
				sum += int(row[clampInt(k, 0, w-1)])
			}
			for x := 0; x < w; x++ {
				out[x] = uint8((sum + window/2) / window)
				sum -= int(row[clampInt(x-radius, 0, w-1)])
				sum += int(row[clampInt(x+radius+1, 0, w-1)])
			}
		}
	})
}

// blurVertical is one vertical box-blur pass with a sliding window. Columns are
// independent, so they are split into chunks just like rows.
func blurVertical(src, dst []uint8, w, h, radius int) {
	window := 2*radius + 1
	parallelChunks(w, func(from, to int) {
		for x := from; x < to; x++ {
			var sum int
			for k := -radius; k <= radius; k++ {
				sum += int(src[clampInt(k, 0, h-1)*w+x])
			}
			for y := 0; y < h; y++ {
				dst[y*w+x] = uint8((sum + window/2) / window)
				sum -= int(src[clampInt(y-radius, 0, h-1)*w+x])
				sum += int(src[clampInt(y+radius+1, 0, h-1)*w+x])
			}
		}
	})
}

// addLumaDelta adds the same delta to the three color channels of one pixel.
// Transparent pixels are left untouched by the callers.
func addLumaDelta(row []uint8, i int, delta float64) {
	row[i] = clampChannel(float64(row[i]) + delta)
	row[i+1] = clampChannel(float64(row[i+1]) + delta)
	row[i+2] = clampChannel(float64(row[i+2]) + delta)
}

// clampChannel rounds and clamps a channel value to 0..255.
func clampChannel(v float64) uint8 {
	return uint8(clampFloat(v+0.5, 0, 255))
}

// applyTexture is an unsharp mask with a small radius: it boosts (or, for a
// negative amount, smooths) the high-frequency detail of the luminance.
func applyTexture(dst *image.NRGBA, amount float64) {
	w, h := dst.Rect.Dx(), dst.Rect.Dy()
	if w == 0 || h == 0 {
		return
	}
	y := lumaPlane(dst)
	radius := int(math.Round(0.0015 * float64(min(w, h))))
	if radius < 1 {
		radius = 1
	}
	blur := boxBlur3(y, w, h, radius)
	k := amount / 100
	parallelChunks(h, func(from, to int) {
		for row := from; row < to; row++ {
			pix := dst.Pix[row*dst.Stride : row*dst.Stride+w*4]
			base := row * w
			for x := 0; x < w; x++ {
				i := x * 4
				if pix[i+3] == 0 {
					continue
				}
				delta := k * (float64(y[base+x]) - float64(blur[base+x]))
				addLumaDelta(pix, i, delta)
			}
		}
	})
}

// applyClarity is an unsharp mask with a large radius. The correction is
// limited to +-40 levels and weighted by the midtone mask, which keeps the
// edges of contrasty objects from ringing.
func applyClarity(dst *image.NRGBA, amount float64) {
	w, h := dst.Rect.Dx(), dst.Rect.Dy()
	if w == 0 || h == 0 {
		return
	}
	y := lumaPlane(dst)
	radius := int(math.Round(0.008 * float64(min(w, h))))
	if radius < 2 {
		radius = 2
	}
	blur := boxBlur3(y, w, h, radius)
	k := amount / 100 * 1.25
	parallelChunks(h, func(from, to int) {
		for row := from; row < to; row++ {
			pix := dst.Pix[row*dst.Stride : row*dst.Stride+w*4]
			base := row * w
			for x := 0; x < w; x++ {
				i := x * 4
				if pix[i+3] == 0 {
					continue
				}
				yVal := float64(y[base+x])
				d := clampFloat(yVal-float64(blur[base+x]), -40, 40)
				t := math.Abs(2*yVal/255 - 1)
				weight := 1 - 0.6*t*t
				addLumaDelta(pix, i, k*d*weight)
			}
		}
	})
}

// applyDehaze removes (positive amount) or adds (negative amount) haze with the
// dark channel prior. The atmosphere and the transmission map are estimated on
// a reduced copy, then the map is interpolated back to the full frame.
func applyDehaze(dst *image.NRGBA, amount float64) *image.NRGBA {
	w, h := dst.Rect.Dx(), dst.Rect.Dy()
	if w == 0 || h == 0 {
		return dst
	}
	small := dst
	sw, sh := w, h
	if max(w, h) > maxDehazeSide {
		scale := float64(maxDehazeSide) / float64(max(w, h))
		sw = max(1, int(math.Round(float64(w)*scale)))
		sh = max(1, int(math.Round(float64(h)*scale)))
		small = imaging.Resize(dst, sw, sh, imaging.Linear)
	}
	sw, sh = small.Rect.Dx(), small.Rect.Dy()

	dark := darkChannel(small)
	atmosphere := atmosphericLight(small, dark)

	tMap := make([]uint8, sw*sh)
	for y := 0; y < sh; y++ {
		row := small.Pix[y*small.Stride : y*small.Stride+sw*4]
		for x := 0; x < sw; x++ {
			i := x * 4
			transmission := 1 - dehazeOmega*min(
				float64(row[i])/atmosphere[0],
				min(float64(row[i+1])/atmosphere[1], float64(row[i+2])/atmosphere[2]))
			transmission = clampFloat(transmission, dehazeMinT, 1)
			tMap[y*sw+x] = uint8(transmission*255 + 0.5)
		}
	}
	tRadius := int(math.Round(0.02 * float64(min(sw, sh))))
	if tRadius < 1 {
		tRadius = 1
	}
	tMap = boxBlur3(tMap, sw, sh, tRadius)

	fullT := tMap
	if sw != w || sh != h {
		gray := image.NewGray(image.Rect(0, 0, sw, sh))
		copy(gray.Pix, tMap)
		fullT = resizeTransmission(gray, w, h)
	}

	k := amount / 100
	if amount > 0 {
		parallelChunks(h, func(from, to int) {
			for row := from; row < to; row++ {
				pix := dst.Pix[row*dst.Stride : row*dst.Stride+w*4]
				base := row * w
				for x := 0; x < w; x++ {
					i := x * 4
					if pix[i+3] == 0 {
						continue
					}
					t := float64(fullT[base+x]) / 255
					if t < dehazeMinT {
						t = dehazeMinT
					}
					for c := 0; c < 3; c++ {
						in := float64(pix[i+c])
						clear := (in-atmosphere[c])/t + atmosphere[c]
						pix[i+c] = clampChannel(in + k*(clear-in))
					}
				}
			}
		})
		return dst
	}

	// Negative amount adds haze: the frame moves toward the atmosphere color
	// and loses contrast.
	haze := -k
	parallelChunks(h, func(from, to int) {
		for row := from; row < to; row++ {
			pix := dst.Pix[row*dst.Stride : row*dst.Stride+w*4]
			for x := 0; x < w; x++ {
				i := x * 4
				if pix[i+3] == 0 {
					continue
				}
				for c := 0; c < 3; c++ {
					in := float64(pix[i+c])
					pix[i+c] = clampChannel(in*(1-0.75*haze) + atmosphere[c]*0.75*haze)
				}
			}
		}
	})
	return dst
}

// darkChannel returns the per-pixel minimum of the RGB channels.
func darkChannel(img *image.NRGBA) []uint8 {
	w, h := img.Rect.Dx(), img.Rect.Dy()
	out := make([]uint8, w*h)
	parallelChunks(h, func(from, to int) {
		for y := from; y < to; y++ {
			row := img.Pix[y*img.Stride : y*img.Stride+w*4]
			dst := out[y*w : (y+1)*w]
			for x := 0; x < w; x++ {
				i := x * 4
				v := row[i]
				if row[i+1] < v {
					v = row[i+1]
				}
				if row[i+2] < v {
					v = row[i+2]
				}
				dst[x] = v
			}
		}
	})
	return out
}

// atmosphericLight estimates the color of the haze: the average RGB of the
// pixels whose dark channel is above the 99.9th percentile, at least one pixel
// and never below one per channel.
func atmosphericLight(img *image.NRGBA, dark []uint8) [3]float64 {
	var atmosphere [3]float64
	if len(dark) == 0 {
		return [3]float64{1, 1, 1}
	}
	var hist [256]int
	maxDark := 0
	for _, v := range dark {
		hist[v]++
		if int(v) > maxDark {
			maxDark = int(v)
		}
	}
	target := int(math.Ceil(0.999 * float64(len(dark))))
	if target < 1 {
		target = 1
	}
	threshold := maxDark
	cumulative := 0
	for v := 0; v < 256; v++ {
		cumulative += hist[v]
		if cumulative >= target {
			threshold = v
			break
		}
	}

	w, h := img.Rect.Dx(), img.Rect.Dy()
	var sum [3]float64
	count := 0
	add := func(x, y int) {
		i := y*img.Stride + x*4
		sum[0] += float64(img.Pix[i])
		sum[1] += float64(img.Pix[i+1])
		sum[2] += float64(img.Pix[i+2])
		count++
	}
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			if int(dark[y*w+x]) > threshold {
				add(x, y)
			}
		}
	}
	if count == 0 {
		// A uniform dark channel has no pixel strictly above the threshold:
		// the brightest pixel still is the best atmosphere estimate.
		found := false
		for y := 0; y < h && !found; y++ {
			for x := 0; x < w; x++ {
				if int(dark[y*w+x]) == maxDark {
					add(x, y)
					found = true
					break
				}
			}
		}
	}
	for c := 0; c < 3; c++ {
		atmosphere[c] = 1
		if count > 0 {
			atmosphere[c] = math.Max(1, sum[c]/float64(count))
		}
	}
	return atmosphere
}

// resizeTransmission upsamples the gray transmission map to the full frame.
// imaging.Resize returns an NRGBA image, so the gray value is the first byte of
// every pixel; slicing its Pix by w*h would read only the first quarter of the
// frame, hence the explicit per-pixel copy.
func resizeTransmission(gray *image.Gray, w, h int) []uint8 {
	resized := imaging.Resize(gray, w, h, imaging.Linear)
	out := make([]uint8, w*h)
	for i := range out {
		out[i] = resized.Pix[i*4]
	}
	return out
}

// applyVignette darkens, brightens, tints or fills the edges of the frame. The
// geometry is normalized to the half-frame: roundness turns the ellipse into a
// rounded rectangle, midpoint sets where the zone starts and feather sets how
// soft the transition is.
func applyVignette(dst *image.NRGBA, v VignetteParams) {
	w, h := dst.Rect.Dx(), dst.Rect.Dy()
	if w == 0 || h == 0 || v.Amount == 0 {
		return
	}
	exponent := 2 * math.Pow(6, v.Roundness/100)
	invExponent := 1 / exponent
	start := 0.15 + 0.85*v.Midpoint/100
	featherWidth := 0.02 + 0.48*v.Feather/100

	xTerms := make([]float64, w)
	for x := 0; x < w; x++ {
		dx := (float64(x)+0.5)/float64(w)*2 - 1
		xTerms[x] = math.Pow(math.Abs(dx), exponent)
	}
	yTerms := make([]float64, h)
	for y := 0; y < h; y++ {
		dy := (float64(y)+0.5)/float64(h)*2 - 1
		yTerms[y] = math.Pow(math.Abs(dy), exponent)
	}

	amount := v.Amount / 100
	tintR, tintG, tintB := parseHexRGB(v.Color)
	tintY := luminance(tintR, tintG, tintB)
	highlights := v.Highlights / 100

	mode := v.Mode
	if mode != VignetteColor && mode != VignetteColorize {
		mode = VignetteLight
	}

	parallelChunks(h, func(from, to int) {
		for row := from; row < to; row++ {
			pix := dst.Pix[row*dst.Stride : row*dst.Stride+w*4]
			termY := yTerms[row]
			for x := 0; x < w; x++ {
				i := x * 4
				if pix[i+3] == 0 {
					continue
				}
				distance := math.Pow(xTerms[x]+termY, invExponent)
				mask := smoothstep(start, start+featherWidth, distance)
				if mask == 0 {
					continue
				}
				yVal := luminance(float64(pix[i]), float64(pix[i+1]), float64(pix[i+2]))
				highlightFactor := highlights * smoothstep(0.5, 1, yVal/255)
				switch mode {
				case VignetteColor:
					strength := math.Abs(amount) * mask * (1 - highlightFactor)
					if strength == 0 {
						continue
					}
					targetR := tintR - tintY + yVal
					targetG := tintG - tintY + yVal
					targetB := tintB - tintY + yVal
					pix[i] = clampChannel(float64(pix[i]) + (targetR-float64(pix[i]))*strength)
					pix[i+1] = clampChannel(float64(pix[i+1]) + (targetG-float64(pix[i+1]))*strength)
					pix[i+2] = clampChannel(float64(pix[i+2]) + (targetB-float64(pix[i+2]))*strength)
				case VignetteColorize:
					strength := math.Abs(amount) * mask * (1 - highlightFactor)
					if strength == 0 {
						continue
					}
					pix[i] = clampChannel(float64(pix[i]) + (tintR-float64(pix[i]))*strength)
					pix[i+1] = clampChannel(float64(pix[i+1]) + (tintG-float64(pix[i+1]))*strength)
					pix[i+2] = clampChannel(float64(pix[i+2]) + (tintB-float64(pix[i+2]))*strength)
				default:
					factor := 1 - amount*0.85*mask
					factor += (1 - factor) * highlightFactor * mask
					pix[i] = clampChannel(float64(pix[i]) * factor)
					pix[i+1] = clampChannel(float64(pix[i+1]) * factor)
					pix[i+2] = clampChannel(float64(pix[i+2]) * factor)
				}
			}
		}
	})
}

// smoothstep is the Hermite ramp between the two edges, clamped to 0..1.
func smoothstep(edge0, edge1, x float64) float64 {
	if edge1 <= edge0 {
		if x < edge1 {
			return 0
		}
		return 1
	}
	t := clamp01((x - edge0) / (edge1 - edge0))
	return t * t * (3 - 2*t)
}

// parseHexRGB decodes a "#rrggbb" color. An unusable color reads as black.
func parseHexRGB(color string) (float64, float64, float64) {
	if len(color) != 7 || color[0] != '#' {
		return 0, 0, 0
	}
	var rgb [3]float64
	for c := 0; c < 3; c++ {
		hi := hexNibble(color[1+c*2])
		lo := hexNibble(color[2+c*2])
		if hi < 0 || lo < 0 {
			return 0, 0, 0
		}
		rgb[c] = float64(hi*16 + lo)
	}
	return rgb[0], rgb[1], rgb[2]
}

// hexNibble decodes one hexadecimal digit, or -1 for an invalid character.
func hexNibble(c byte) int {
	switch {
	case c >= '0' && c <= '9':
		return int(c - '0')
	case c >= 'a' && c <= 'f':
		return int(c-'a') + 10
	case c >= 'A' && c <= 'F':
		return int(c-'A') + 10
	}
	return -1
}

// Grain noise is a pure function of the pixel coordinates and the seed. Using
// the splitmix64 finalizer as a hash (instead of a stateful sequence) keeps the
// noise independent of the row order, so the row-parallel render is bit-for-bit
// reproducible.

const (
	grainSaltValue uint64 = 0x9e3779b97f4a7c15
	grainSaltWhite uint64 = 0xd1b54a32d192ed03
)

// mix64 is the splitmix64 finalizer.
func mix64(x uint64) uint64 {
	x += 0x9e3779b97f4a7c15
	x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9
	x = (x ^ (x >> 27)) * 0x94d049bb133111eb
	return x ^ (x >> 31)
}

// hashSigned maps a seed, a salt and integer coordinates to [-1, 1].
func hashSigned(seed, salt uint64, x, y int) float64 {
	z := seed ^ salt
	z ^= uint64(x) * 0x9e3779b97f4a7c15
	z ^= uint64(y) * 0xc2b2ae3d27d4eb4f
	h := mix64(z)
	return float64(h>>11)/float64(uint64(1)<<53)*2 - 1
}

// applyGrain adds a deterministic luminance noise. Size sets the cell of the
// smooth value noise and roughness mixes it with the per-pixel white noise; the
// weight fades the noise out of the darkest and brightest tones.
func applyGrain(dst *image.NRGBA, g GrainParams, seed uint64) {
	w, h := dst.Rect.Dx(), dst.Rect.Dy()
	if w == 0 || h == 0 || g.Amount == 0 {
		return
	}
	cell := int(math.Round((0.0005 + 0.004*g.Size/100) * float64(min(w, h))))
	if cell < 1 {
		cell = 1
	}

	gridW := (w-1)/cell + 2
	gridH := (h-1)/cell + 2
	grid := make([]float64, gridW*gridH)
	for gy := 0; gy < gridH; gy++ {
		for gx := 0; gx < gridW; gx++ {
			grid[gy*gridW+gx] = hashSigned(seed, grainSaltValue, gx, gy)
		}
	}

	roughness := g.Roughness / 100
	smoothWeight := 1 - roughness
	strength := g.Amount / 100 * 38

	parallelChunks(h, func(from, to int) {
		for row := from; row < to; row++ {
			pix := dst.Pix[row*dst.Stride : row*dst.Stride+w*4]
			gy := row / cell
			fy := float64(row-gy*cell) / float64(cell)
			vy := smoothstep(0, 1, fy)
			for x := 0; x < w; x++ {
				i := x * 4
				if pix[i+3] == 0 {
					continue
				}
				gx := x / cell
				fx := float64(x-gx*cell) / float64(cell)
				vx := smoothstep(0, 1, fx)

				a := grid[gy*gridW+gx]
				b := grid[gy*gridW+gx+1]
				c := grid[(gy+1)*gridW+gx]
				d := grid[(gy+1)*gridW+gx+1]
				top := a + (b-a)*vx
				bottom := c + (d-c)*vx
				value := top + (bottom-top)*vy

				white := hashSigned(seed, grainSaltWhite, x, row)
				noise := value*smoothWeight + white*roughness

				yVal := luminance(float64(pix[i]), float64(pix[i+1]), float64(pix[i+2]))
				t := 2*yVal/255 - 1
				weight := 1 - t*t
				addLumaDelta(pix, i, noise*strength*weight)
			}
		}
	})
}