package main

import (
	"bytes"
	"errors"
	"image"
	"image/color"
	"image/draw"
	"image/png"
	"math"
	"os/exec"

	"github.com/liyue201/goqr"
	"github.com/makiuchi-d/gozxing"
	qrx "github.com/makiuchi-d/gozxing/qrcode"
	qrgen "github.com/skip2/go-qrcode"
)

// GenerateQRPNG returns PNG bytes for the given text.
func GenerateQRPNG(text string, size int) ([]byte, error) {
	if size <= 0 {
		size = 256
	}
	// Use Low EC for larger modules; improves decode reliability for long payloads
	png, err := qrgen.Encode(text, qrgen.Low, size)
	if err != nil {
		return nil, err
	}
	return png, nil
}

// DecodeQR decodes first QR code text from an image.
func DecodeQR(img image.Image) (string, error) {
	// Try basic decode first
	codes, err := goqr.Recognize(img)
	if err == nil && len(codes) > 0 {
		return string(codes[0].Payload), nil
	}
	if err != nil {
	}

	// Convert palette/indexed images to RGBA first to avoid issues with goqr
	if _, ok := img.ColorModel().(color.Palette); ok {
		bounds := img.Bounds()
		rgba := image.NewRGBA(bounds)
		draw.Draw(rgba, bounds, img, bounds.Min, draw.Src)
		img = rgba

		// Try again after conversion
		codes, err := goqr.Recognize(img)
		if err == nil && len(codes) > 0 {
			return string(codes[0].Payload), nil
		}
	}

	// Try multiple preprocessing strategies to improve recognition from clipboard screenshots
	attempts := []image.Image{}
	add := func(im image.Image) {
		if im != nil {
			attempts = append(attempts, im)
		}
	}
	add(img)

	bounds := img.Bounds()
	w, h := bounds.Dx(), bounds.Dy()

	// Helper: nearest-neighbour scale
	scale := func(src image.Image, fw, fh int) image.Image {
		if fw <= 0 || fh <= 0 {
			return src
		}
		dst := image.NewNRGBA(image.Rect(0, 0, fw, fh))
		sb := src.Bounds()
		sw, sh := float64(sb.Dx()), float64(sb.Dy())
		for y := 0; y < fh; y++ {
			sy := int(float64(y)/float64(fh)*sh + 0.5)
			if sy >= sb.Dy() {
				sy = sb.Dy() - 1
			}
			for x := 0; x < fw; x++ {
				sx := int(float64(x)/float64(fw)*sw + 0.5)
				if sx >= sb.Dx() {
					sx = sb.Dx() - 1
				}
				dst.Set(x, y, src.At(sb.Min.X+sx, sb.Min.Y+sy))
			}
		}
		return dst
	}

	// Helper: binarize with adaptive threshold (simple mean + variance fudge)
	binarize := func(src image.Image) image.Image {
		b := src.Bounds()
		gray := image.NewGray(b)
		var sum, sum2 float64
		total := float64(b.Dx() * b.Dy())
		for y := b.Min.Y; y < b.Max.Y; y++ {
			for x := b.Min.X; x < b.Max.X; x++ {
				r, g, bl, _ := src.At(x, y).RGBA()
				lum := float64((r*299 + g*587 + bl*114) / 1000 >> 8) // 0-255 approx
				sum += lum
				sum2 += lum * lum
			}
		}
		mean := sum / total
		variance := (sum2 / total) - mean*mean
		stddev := math.Sqrt(math.Max(variance, 0))
		// threshold tweaks: lower slightly for dark backgrounds
		thr := mean - stddev*0.15
		if thr < 50 {
			thr = 50
		}
		if thr > 200 {
			thr = 200
		}
		for y := b.Min.Y; y < b.Max.Y; y++ {
			for x := b.Min.X; x < b.Max.X; x++ {
				r, g, bl, _ := src.At(x, y).RGBA()
				lum := float64((r*299 + g*587 + bl*114) / 1000 >> 8)
				var c color.Gray
				if lum < thr {
					c = color.Gray{Y: 0}
				} else {
					c = color.Gray{Y: 255}
				}
				gray.Set(x, y, c)
			}
		}
		return gray
	}

	// Scale strategy: even if image is not tiny, scaled variants sometimes help goqr lock onto module grid
	// 1) If very small, aggressive scale
	if w < 280 || h < 280 {
		factor := 2
		if w < 160 {
			factor = 3
		}
		if w < 100 {
			factor = 4
		}
		add(scale(img, w*factor, h*factor))
	}
	// 2) Always try moderate upscale variants for mid-sized screenshots (helps when original was downscaled with interpolation)
	if w >= 180 && w <= 900 { // avoid exploding huge images
		add(scale(img, w*2, h*2))
		if w*3 < 2000 && h*3 < 2000 { // guard upper bound
			add(scale(img, w*3, h*3))
		}
	}

	// Add binarized versions (original + scaled)
	add(binarize(img))
	if w < 320 || h < 320 {
		factor := 2
		add(binarize(scale(img, w*factor, h*factor)))
	}
	// Also binarize any newly added large upscales (skip if too big)
	for _, cand := range attempts { // safe: attempts already contains originals
		cw := cand.Bounds().Dx()
		ch := cand.Bounds().Dy()
		if (cw <= 1400 && ch <= 1400) && (cw > w || ch > h) { // only for upscaled variants within limit
			add(binarize(cand))
		}
	}

	// Add contrast-enhanced alpha-flattened variant
	flatten := func(src image.Image) image.Image {
		b := src.Bounds()
		dst := image.NewRGBA(b)
		draw.Draw(dst, b, image.Black, image.Point{}, draw.Src)
		draw.Draw(dst, b, src, b.Min, draw.Over)
		return dst
	}
	add(flatten(img))

	// Rotations (90,180,270) in case of orientation issues
	rotate := func(src image.Image) []image.Image {
		b := src.Bounds()
		w := b.Dx()
		h := b.Dy()
		r90 := image.NewRGBA(image.Rect(0, 0, h, w))
		r180 := image.NewRGBA(image.Rect(0, 0, w, h))
		r270 := image.NewRGBA(image.Rect(0, 0, h, w))
		for y := 0; y < h; y++ {
			for x := 0; x < w; x++ {
				c := src.At(b.Min.X+x, b.Min.Y+y)
				// 90
				r90.Set(h-1-y, x, c)
				// 180
				r180.Set(w-1-x, h-1-y, c)
				// 270
				r270.Set(y, w-1-x, c)
			}
		}
		return []image.Image{r90, r180, r270}
	}
	for _, r := range rotate(img) {
		add(r)
	}

	var firstErr error
	for _, attempt := range attempts {
		codes, err := goqr.Recognize(attempt)
		if err != nil {
			if firstErr == nil {
				firstErr = err
			}
			continue
		}
		if len(codes) > 0 {
			return string(codes[0].Payload), nil
		}
	}

	// ZXing helper used across multiple branches
	decodeZX := func(src image.Image, hints map[gozxing.DecodeHintType]interface{}) (string, error) {
		binSrc := gozxing.NewLuminanceSourceFromImage(src)
		// Try hybrid first (usually better for photos)
		bmpH, _ := gozxing.NewBinaryBitmap(gozxing.NewHybridBinarizer(binSrc))
		reader := qrx.NewQRCodeReader()
		if res, err := reader.Decode(bmpH, hints); err == nil {
			return res.GetText(), nil
		}
		// Fallback to global histogram
		bmpG, _ := gozxing.NewBinaryBitmap(gozxing.NewGlobalHistgramBinarizer(binSrc))
		if res, err := reader.Decode(bmpG, hints); err == nil {
			return res.GetText(), nil
		} else {
			return "", err
		}
	}
	// Also try ZXing on all preprocessed attempts
	for _, attempt := range attempts {
		if txt, err := decodeZX(attempt, map[gozxing.DecodeHintType]interface{}{
			gozxing.DecodeHintType_TRY_HARDER:   true,
			gozxing.DecodeHintType_PURE_BARCODE: true,
		}); err == nil && txt != "" {
			return txt, nil
		}
		if txt, err := decodeZX(attempt, map[gozxing.DecodeHintType]interface{}{
			gozxing.DecodeHintType_TRY_HARDER: true,
		}); err == nil && txt != "" {
			return txt, nil
		}
	}
	// --- Heuristic fallback: auto-invert + quiet-zone pad + gozxing ---
	addQuietZone := func(src image.Image) image.Image {
		b := src.Bounds()
		pad := int(math.Max(float64(b.Dx()/32), 8))
		out := image.NewRGBA(image.Rect(0, 0, b.Dx()+pad*2, b.Dy()+pad*2))
		// fill white
		for y := 0; y < out.Bounds().Dy(); y++ {
			for x := 0; x < out.Bounds().Dx(); x++ {
				out.Set(x, y, color.White)
			}
		}
		draw.Draw(out, b.Add(image.Point{X: pad, Y: pad}), src, b.Min, draw.Src)
		return out
	}
	invert := func(src image.Image) image.Image {
		b := src.Bounds()
		out := image.NewRGBA(b)
		for y := b.Min.Y; y < b.Max.Y; y++ {
			for x := b.Min.X; x < b.Max.X; x++ {
				r, g, bb, a := src.At(x, y).RGBA()
				out.Set(x, y, color.RGBA{uint8(255 - r/257), uint8(255 - g/257), uint8(255 - bb/257), uint8(a / 257)})
			}
		}
		return out
	}
	// estimate dark pixel ratio
	b := img.Bounds()
	var dark, total int
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			r, g, bb, _ := img.At(x, y).RGBA()
			lum := (r*299 + g*587 + bb*114) / 1000
			if lum < 128<<8 {
				dark++
			}
			total++
		}
	}
	ratio := float64(dark) / float64(total)
	needInvert := ratio > 0.85 // almost full dark background (white modules?)
	var modified []image.Image
	base := img
	if needInvert {
		base = invert(base)
	}
	modified = append(modified, addQuietZone(base))
	modified = append(modified, addQuietZone(invert(base)))
	for _, m := range modified {
		if codes, err := goqr.Recognize(m); err == nil && len(codes) > 0 {
			return string(codes[0].Payload), nil
		}
		// Also try ZXing on these variants before moving on
		if txt, err := decodeZX(m, map[gozxing.DecodeHintType]interface{}{
			gozxing.DecodeHintType_TRY_HARDER:   true,
			gozxing.DecodeHintType_PURE_BARCODE: true,
		}); err == nil && txt != "" {
			return txt, nil
		}
		if txt, err := decodeZX(m, map[gozxing.DecodeHintType]interface{}{
			gozxing.DecodeHintType_TRY_HARDER: true,
		}); err == nil && txt != "" {
			return txt, nil
		}
	}
	// ZXing with TRY_HARDER and PURE_BARCODE hints which often fix dense PNGs
	hints := map[gozxing.DecodeHintType]interface{}{
		gozxing.DecodeHintType_TRY_HARDER: true,
		// Many of our images are perfect render PNGs from go-qrcode -> enable PURE_BARCODE
		gozxing.DecodeHintType_PURE_BARCODE: true,
	}
	if txt, err := decodeZX(base, hints); err == nil && txt != "" {
		return txt, nil
	}
	// Try again without PURE_BARCODE in case quiet zone is cropped
	hints2 := map[gozxing.DecodeHintType]interface{}{gozxing.DecodeHintType_TRY_HARDER: true}
	if txt, err := decodeZX(base, hints2); err == nil && txt != "" {
		return txt, nil
	}
	// Try ZXing on rotations as well
	for _, r := range rotate(base) {
		if txt, err := decodeZX(r, hints2); err == nil && txt != "" {
			return txt, nil
		}
		if txt, err := decodeZX(r, hints); err == nil && txt != "" {
			return txt, nil
		}
	}

	// Extra pure-QR recovery: try small insets and simple upscales with PURE_BARCODE.
	inset := func(src image.Image, px int) image.Image {
		if px <= 0 {
			return src
		}
		b := src.Bounds()
		r := image.Rect(b.Min.X+px, b.Min.Y+px, b.Max.X-px, b.Max.Y-px)
		if r.Dx() <= 10 || r.Dy() <= 10 {
			return src
		}
		out := image.NewRGBA(image.Rect(0, 0, r.Dx(), r.Dy()))
		draw.Draw(out, out.Bounds(), src, r.Min, draw.Src)
		return out
	}
	cropSides := func(src image.Image, l, t, r, b int) image.Image {
		bb := src.Bounds()
		rr := image.Rect(bb.Min.X+l, bb.Min.Y+t, bb.Max.X-r, bb.Max.Y-b)
		if rr.Dx() <= 10 || rr.Dy() <= 10 || rr.Min.X >= rr.Max.X || rr.Min.Y >= rr.Max.Y {
			return src
		}
		out := image.NewRGBA(image.Rect(0, 0, rr.Dx(), rr.Dy()))
		draw.Draw(out, out.Bounds(), src, rr.Min, draw.Src)
		return out
	}
	tryPure := func(src image.Image) (string, bool) {
		for _, px := range []int{1, 2, 3} {
			if txt, err := decodeZX(inset(src, px), hints); err == nil && txt != "" {
				return txt, true
			}
		}
		return "", false
	}
	// Try on base, binarized base and scaled variants
	if txt, ok := tryPure(base); ok {
		return txt, nil
	}
	if txt, ok := tryPure(binarize(base)); ok {
		return txt, nil
	}
	if w < 800 && h < 800 { // upscale to mitigate rounding of module size
		if txt, ok := tryPure(scale(base, w*2, h*2)); ok {
			return txt, nil
		}
		if txt, ok := tryPure(scale(base, w*3, h*3)); ok {
			return txt, nil
		}
	}

	// Brute-force small asymmetric crops to satisfy ZXing's 1 (mod 4) dimension constraint.
	for _, src := range []image.Image{base, binarize(base)} {
		for l := 0; l <= 2; l++ {
			for t := 0; t <= 2; t++ {
				for r := 0; r <= 2; r++ {
					for btm := 0; btm <= 2; btm++ {
						cand := cropSides(src, l, t, r, btm)
						if cand == src {
							continue
						}
						if txt, err := decodeZX(cand, hints); err == nil && txt != "" {
							return txt, nil
						}
						if txt, err := decodeZX(cand, hints2); err == nil && txt != "" {
							return txt, nil
						}
					}
				}
			}
		}
	}

	// Bounding box crop of dark pixels, then upscale and retry
	boundsAll := base.Bounds()
	minX, minY := boundsAll.Max.X, boundsAll.Max.Y
	maxX, maxY := boundsAll.Min.X, boundsAll.Min.Y
	for y := boundsAll.Min.Y; y < boundsAll.Max.Y; y++ {
		for x := boundsAll.Min.X; x < boundsAll.Max.X; x++ {
			r, g, bb, _ := base.At(x, y).RGBA()
			lum := (r*299 + g*587 + bb*114) / 1000
			if lum < 180<<8 { // treat as dark
				if x < minX {
					minX = x
				}
				if y < minY {
					minY = y
				}
				if x > maxX {
					maxX = x
				}
				if y > maxY {
					maxY = y
				}
			}
		}
	}
	if minX < maxX && minY < maxY {
		crop := image.NewRGBA(image.Rect(0, 0, maxX-minX+1, maxY-minY+1))
		draw.Draw(crop, crop.Bounds(), base, image.Point{X: minX, Y: minY}, draw.Src)
		for _, up := range []int{2, 3, 4} {
			cw := crop.Bounds().Dx() * up
			ch := crop.Bounds().Dy() * up
			if cw > 2400 || ch > 2400 {
				continue
			}
			scaled := image.NewRGBA(image.Rect(0, 0, cw, ch))
			for y := 0; y < ch; y++ {
				sy := y / up
				for x := 0; x < cw; x++ {
					sx := x / up
					scaled.Set(x, y, crop.At(sx, sy))
				}
			}
			if txt, err := decodeZX(scaled, hints2); err == nil && txt != "" {
				return txt, nil
			}
		}
	}
	if firstErr != nil {
		return "", firstErr
	}
	// Final external fallback: try zbarimg if available on the system
	if cmd, lookErr := exec.LookPath("zbarimg"); lookErr == nil {
		buf := &bytes.Buffer{}
		if err := png.Encode(buf, img); err == nil {
			c := exec.Command(cmd, "-q", "--raw", "-")
			c.Stdin = bytes.NewReader(buf.Bytes())
			out, err := c.Output()
			if err == nil && len(out) > 0 {
				return string(out), nil
			}
		}
	}
	return "", errors.New("no qr code found")
}

// LoadPNG decodes raw PNG bytes to image.Image.
func LoadPNG(b []byte) (image.Image, error) {
	im, err := png.Decode(bytes.NewReader(b))
	if err != nil {
		return nil, err
	}
	return im, nil
}