package onnx

// Ported from meikiocr (https://github.com/rtr46/meikiocr), Apache-2.0.
// Model weights: rtr46/meiki.txt.recognition.v0, LGPL-3.0.

import (
	"fmt"
	"image"
	"math"
	"sort"
	"strings"

	ort "github.com/yalue/onnxruntime_go"
	"golang.org/x/image/draw"
)

const (
	recInputHeight   = 32
	recInputWidth    = 960
	defaultRecConf   = 0.1
	recOverlapThresh = 0.3
)

var swappedPairs = map[string]string{
	"儡傀": "傀儡",
	"談冗": "冗談",
	"汰淘": "淘汰",
	"沱滂": "滂沱",
	"攣痙": "痙攣",
	"酊酩": "酩酊",
	"麭麺": "麺麭",
	"哭慟": "慟哭",
}

// MeikiRecognizer runs the horizontal character recognition model.
type MeikiRecognizer struct {
	sess *ort.DynamicAdvancedSession
}

// NewMeikiRecognizer loads the horizontal text recognizer ONNX model.
func NewMeikiRecognizer(modelPath string) (*MeikiRecognizer, error) {
	sess, err := newSession(modelPath, []string{"images", "orig_target_sizes"}, []string{"char_codes", "boxes", "scores"})
	if err != nil {
		return nil, fmt.Errorf("load meiki recognizer: %w", err)
	}
	return &MeikiRecognizer{sess: sess}, nil
}

// Close destroys the underlying ORT session.
func (r *MeikiRecognizer) Close() error {
	if r.sess != nil {
		err := r.sess.Destroy()
		r.sess = nil
		return err
	}
	return nil
}

type charCandidate struct {
	char     rune
	score    float32
	interval [2]float64
}

// RecognizeLine recognizes one horizontal cropped text line image.
func (r *MeikiRecognizer) RecognizeLine(crop image.Image) (string, error) {
	bounds := crop.Bounds()
	cropW := bounds.Dx()
	cropH := bounds.Dy()
	if cropW <= 0 || cropH <= 0 {
		return "", nil
	}

	scale := float64(recInputHeight) / float64(cropH)
	newW := int(math.Round(float64(cropW) * scale))
	newH := recInputHeight

	if newW > recInputWidth {
		scaleW := float64(recInputWidth) / float64(newW)
		newW = recInputWidth
		newH = int(math.Round(float64(recInputHeight) * scaleW))
	}
	if newW <= 0 || newH <= 0 {
		return "", nil
	}

	resized := image.NewRGBA(image.Rect(0, 0, newW, newH))
	draw.BiLinear.Scale(resized, resized.Bounds(), crop, bounds, draw.Over, nil)

	// Prepare tensor [1, 3, 32, 960] in BGR layout
	imgData := make([]float32, 1*3*recInputHeight*recInputWidth)
	strideH := recInputWidth
	strideC := recInputHeight * recInputWidth

	for y := 0; y < newH; y++ {
		for x := 0; x < newW; x++ {
			c := resized.RGBAAt(x, y)
			pxIdx := y*strideH + x
			imgData[0*strideC+pxIdx] = float32(c.B) / 255.0
			imgData[1*strideC+pxIdx] = float32(c.G) / 255.0
			imgData[2*strideC+pxIdx] = float32(c.R) / 255.0
		}
	}

	imgTensor, err := ort.NewTensor(ort.Shape{1, 3, recInputHeight, recInputWidth}, imgData)
	if err != nil {
		return "", fmt.Errorf("create rec images tensor: %w", err)
	}
	defer imgTensor.Destroy()

	sizesData := []int64{recInputWidth, recInputHeight}
	sizesTensor, err := ort.NewTensor(ort.Shape{1, 2}, sizesData)
	if err != nil {
		return "", fmt.Errorf("create rec orig_target_sizes tensor: %w", err)
	}
	defer sizesTensor.Destroy()

	outputs := make([]ort.Value, 3)
	err = r.sess.Run([]ort.Value{imgTensor, sizesTensor}, outputs)
	if err != nil {
		return "", fmt.Errorf("meiki rec inference: %w", err)
	}
	defer func() {
		for _, o := range outputs {
			if o != nil {
				o.Destroy()
			}
		}
	}()

	charCodesTensor := outputs[0].(*ort.Tensor[int32])
	boxesTensor := outputs[1].(*ort.Tensor[float32])
	scoresTensor := outputs[2].(*ort.Tensor[float32])

	charCodes := charCodesTensor.GetData()
	boxes := boxesTensor.GetData()
	scores := scoresTensor.GetData()

	var candidates []charCandidate
	for i, score := range scores {
		if score < defaultRecConf {
			continue
		}
		ch := rune(charCodes[i])
		bOffset := i * 4
		rx1 := float64(boxes[bOffset])
		rx2 := float64(boxes[bOffset+2])

		if rx1 >= float64(newW) {
			continue
		}
		rx1 = math.Min(rx1, float64(newW))
		rx2 = math.Min(rx2, float64(newW))
		cx1 := (rx1 / float64(newW)) * float64(cropW)
		cx2 := (rx2 / float64(newW)) * float64(cropW)

		candidates = append(candidates, charCandidate{
			char:     ch,
			score:    score,
			interval: [2]float64{cx1, cx2},
		})
	}

	// 1D NMS: sort candidates descending by score
	sort.Slice(candidates, func(i, j int) bool {
		return candidates[i].score > candidates[j].score
	})

	var accepted []charCandidate
	const epsilon = 1e-6

	for _, cand := range candidates {
		i1C, i2C := cand.interval[0], cand.interval[1]
		lenC := i2C - i1C + epsilon
		isOverlap := false

		for _, acc := range accepted {
			i1A, i2A := acc.interval[0], acc.interval[1]
			if (i1C >= i2A) || (i1A >= i2C) {
				continue
			}
			interStart := math.Max(i1C, i1A)
			interEnd := math.Min(i2C, i2A)
			interLen := math.Max(0, interEnd-interStart)
			lenA := i2A - i1A + epsilon
			minLen := math.Min(lenC, lenA)

			if (interLen / minLen) > recOverlapThresh {
				isOverlap = true
				break
			}
		}

		if !isOverlap {
			accepted = append(accepted, cand)
		}
	}

	// Sort accepted characters left-to-right along the line
	sort.Slice(accepted, func(i, j int) bool {
		return accepted[i].interval[0] < accepted[j].interval[0]
	})

	var b strings.Builder
	for _, c := range accepted {
		b.WriteRune(c.char)
	}
	text := b.String()

	// Fix swapped pairs from known training corpus artifacts
	for wrong, correct := range swappedPairs {
		text = strings.ReplaceAll(text, wrong, correct)
	}

	return text, nil
}