// Package digitocr 识别固定字体、干净背景的多位数字图（如 4 位橙色数码字验证码）。
// 思路：橙色阈值二值化 → 找前景 bbox → 等分 N 列 → 每块归一到固定网格 → 与 0-9 模板做 Hamming 距离。
package digitocr

import (
	"errors"
	"fmt"
	"image"
	"image/color"
	_ "image/png"
	"io"
	"os"
	"sort"
)

const (
	GridW = 10 // 单数字画布宽（覆盖最宽字符 + 余量）
	GridH = 18 // 单数字画布高
)

// Bitmap 单个数字归一后的位图（按行展开）。
type Bitmap [GridW * GridH]uint8

// Options 控制识别行为。零值即默认 4 位、橙色前景。
type Options struct {
	Digits     int                  // 期望位数，默认 4
	IsForeground func(color.Color) bool // 自定义前景判定，nil 时用 IsOrange
}

// Recognize 从文件路径读图并识别。
func Recognize(path string, opts Options) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()
	return RecognizeReader(f, opts)
}

// RecognizeReader 从 io.Reader 读图并识别。
func RecognizeReader(r io.Reader, opts Options) (string, error) {
	img, _, err := image.Decode(r)
	if err != nil {
		return "", err
	}
	return RecognizeImage(img, opts)
}

// RecognizeImage 对已解码的 image.Image 做识别。
func RecognizeImage(img image.Image, opts Options) (string, error) {
	if opts.Digits == 0 {
		opts.Digits = 4
	}
	if opts.IsForeground == nil {
		opts.IsForeground = IsOrange
	}

	bin, bbox := Binarize(img, opts.IsForeground)
	if bbox.Empty() {
		return "", errors.New("digitocr: no foreground pixels")
	}

	cells := SegmentDigits(bin, bbox, opts.Digits)
	out := make([]byte, len(cells))
	for i, cell := range cells {
		bm := Normalize(bin, cell)
		out[i] = '0' + byte(MatchDigit(bm))
	}
	if len(out) != opts.Digits {
		return string(out), fmt.Errorf("digitocr: expected %d digits, segmented %d", opts.Digits, len(out))
	}
	return string(out), nil
}

// SegmentDigits 优先用 8 连通分量切分，每个数字应是一个 CC。
// CC 数不匹配时退回列投影找零列；再不匹配退回等宽切分。
func SegmentDigits(bin [][]bool, bbox image.Rectangle, want int) []image.Rectangle {
	if rects := ConnectedComponents(bin); len(rects) == want {
		return rects
	}
	if rects := projectionSplit(bin, bbox); len(rects) == want {
		return rects
	}
	// 兜底：等宽切分
	out := make([]image.Rectangle, want)
	cellW := bbox.Dx() / want
	for i := 0; i < want; i++ {
		out[i] = image.Rect(
			bbox.Min.X+i*cellW, bbox.Min.Y,
			bbox.Min.X+(i+1)*cellW, bbox.Max.Y,
		)
		if i == want-1 {
			out[i].Max.X = bbox.Max.X
		}
	}
	return out
}

// ConnectedComponents 返回所有 8 连通前景分量的 bbox，按左上 x 升序排列。
// 小于 minPixels 的噪点分量会被丢弃（默认 2）。
func ConnectedComponents(bin [][]bool) []image.Rectangle {
	const minPixels = 2
	h := len(bin)
	if h == 0 {
		return nil
	}
	w := len(bin[0])
	visited := make([][]bool, h)
	for i := range visited {
		visited[i] = make([]bool, w)
	}
	var rects []image.Rectangle
	stack := make([][2]int, 0, 64)
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			if !bin[y][x] || visited[y][x] {
				continue
			}
			minX, minY := x, y
			maxX, maxY := x, y
			pixCount := 0
			stack = append(stack[:0], [2]int{x, y})
			visited[y][x] = true
			for len(stack) > 0 {
				p := stack[len(stack)-1]
				stack = stack[:len(stack)-1]
				pixCount++
				if p[0] < minX {
					minX = p[0]
				}
				if p[1] < minY {
					minY = p[1]
				}
				if p[0] > maxX {
					maxX = p[0]
				}
				if p[1] > maxY {
					maxY = p[1]
				}
				for dy := -1; dy <= 1; dy++ {
					for dx := -1; dx <= 1; dx++ {
						if dx == 0 && dy == 0 {
							continue
						}
						nx, ny := p[0]+dx, p[1]+dy
						if nx >= 0 && nx < w && ny >= 0 && ny < h && bin[ny][nx] && !visited[ny][nx] {
							visited[ny][nx] = true
							stack = append(stack, [2]int{nx, ny})
						}
					}
				}
			}
			if pixCount >= minPixels {
				rects = append(rects, image.Rect(minX, minY, maxX+1, maxY+1))
			}
		}
	}
	sort.Slice(rects, func(i, j int) bool { return rects[i].Min.X < rects[j].Min.X })
	return rects
}

func projectionSplit(bin [][]bool, bbox image.Rectangle) []image.Rectangle {
	cols := make([]int, bbox.Dx())
	for y := bbox.Min.Y; y < bbox.Max.Y; y++ {
		for x := bbox.Min.X; x < bbox.Max.X; x++ {
			if bin[y][x] {
				cols[x-bbox.Min.X]++
			}
		}
	}
	var rects []image.Rectangle
	inRun, runStart := false, 0
	for i := 0; i <= len(cols); i++ {
		present := i < len(cols) && cols[i] > 0
		if present && !inRun {
			inRun = true
			runStart = i
		} else if !present && inRun {
			inRun = false
			rects = append(rects, image.Rect(
				bbox.Min.X+runStart, bbox.Min.Y,
				bbox.Min.X+i, bbox.Max.Y,
			))
		}
	}
	return rects
}

// IsOrange 经验阈值：识别图中那种橙色前景像素。
// 如背景色调差异较大，可换 HSV 判 H∈[10°,30°]。
func IsOrange(c color.Color) bool {
	r, g, b, _ := c.RGBA()
	r8, g8, b8 := r>>8, g>>8, b>>8
	return r8 > 200 && g8 > 80 && g8 < 180 && b8 < 100
}

// Binarize 把图像转成 bool 矩阵 + 所有前景像素的边界框。
// 默认会通过 StripFrame 去掉与图像边缘相连的"外框"前景（如圆角矩形装饰边）。
func Binarize(img image.Image, fg func(color.Color) bool) ([][]bool, image.Rectangle) {
	b := img.Bounds()
	w, h := b.Dx(), b.Dy()
	bin := make([][]bool, h)
	for y := 0; y < h; y++ {
		bin[y] = make([]bool, w)
		for x := 0; x < w; x++ {
			if fg(img.At(b.Min.X+x, b.Min.Y+y)) {
				bin[y][x] = true
			}
		}
	}
	StripFrame(bin)
	return bin, computeBBox(bin)
}

// StripFrame 用 8 邻接 flood fill 从图像四边把任何相连的前景"吃掉"。
// 适用于验证码常见的圆角矩形装饰边——只要数字本身和外框不相连，外框就会被清干净。
func StripFrame(bin [][]bool) {
	h := len(bin)
	if h == 0 {
		return
	}
	w := len(bin[0])
	type pt struct{ x, y int }
	var stack []pt
	push := func(x, y int) {
		if x >= 0 && x < w && y >= 0 && y < h && bin[y][x] {
			bin[y][x] = false
			stack = append(stack, pt{x, y})
		}
	}
	for x := 0; x < w; x++ {
		push(x, 0)
		push(x, h-1)
	}
	for y := 0; y < h; y++ {
		push(0, y)
		push(w-1, y)
	}
	for len(stack) > 0 {
		p := stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		for dy := -1; dy <= 1; dy++ {
			for dx := -1; dx <= 1; dx++ {
				if dx == 0 && dy == 0 {
					continue
				}
				push(p.x+dx, p.y+dy)
			}
		}
	}
}

func computeBBox(bin [][]bool) image.Rectangle {
	h := len(bin)
	if h == 0 {
		return image.Rectangle{}
	}
	w := len(bin[0])
	minX, minY := w, h
	maxX, maxY := -1, -1
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			if bin[y][x] {
				if x < minX {
					minX = x
				}
				if y < minY {
					minY = y
				}
				if x > maxX {
					maxX = x
				}
				if y > maxY {
					maxY = y
				}
			}
		}
	}
	if maxX < 0 {
		return image.Rectangle{}
	}
	return image.Rect(minX, minY, maxX+1, maxY+1)
}

// Normalize 把 cell 内的前景按原尺寸贴到 GridW x GridH 画布的左上角。
// 字体固定时不做缩放采样，模板就是真实像素，区分度最高。
// 若数字尺寸超出画布，多出的右/下部分被截断（GridW/GridH 应当设大于最大字符）。
func Normalize(bin [][]bool, cell image.Rectangle) Bitmap {
	minX, minY := cell.Max.X, cell.Max.Y
	maxX, maxY := cell.Min.X-1, cell.Min.Y-1
	for y := cell.Min.Y; y < cell.Max.Y; y++ {
		if y < 0 || y >= len(bin) {
			continue
		}
		for x := cell.Min.X; x < cell.Max.X; x++ {
			if x < 0 || x >= len(bin[y]) {
				continue
			}
			if bin[y][x] {
				if x < minX {
					minX = x
				}
				if y < minY {
					minY = y
				}
				if x > maxX {
					maxX = x
				}
				if y > maxY {
					maxY = y
				}
			}
		}
	}
	var bm Bitmap
	if minX > maxX {
		return bm
	}
	for y := minY; y <= maxY; y++ {
		gy := y - minY
		if gy >= GridH {
			break
		}
		for x := minX; x <= maxX; x++ {
			gx := x - minX
			if gx >= GridW {
				break
			}
			if bin[y][x] {
				bm[gy*GridW+gx] = 1
			}
		}
	}
	return bm
}

// MatchDigit 与 10 个模板比 Hamming 距离，返回最像的数字。
func MatchDigit(bm Bitmap) int {
	best, bestD := 0, 1<<30
	for d := 0; d < 10; d++ {
		dist := hamming(bm, Templates[d])
		if dist < bestD {
			bestD = dist
			best = d
		}
	}
	return best
}

func hamming(a, b Bitmap) int {
	n := 0
	for i := range a {
		if a[i] != b[i] {
			n++
		}
	}
	return n
}