package onnx

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

import (
	"fmt"
	"image"
	"math"

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

const (
	detInputWidth  = 960
	detInputHeight = 544
	defaultDetConf = 0.5
)

// TextBox is a detected bounding box for a line or column of text.
type TextBox struct {
	X1, Y1, X2, Y2 int
	Score          float32
	IsVertical     bool
}

func (b TextBox) Width() int  { return b.X2 - b.X1 }
func (b TextBox) Height() int { return b.Y2 - b.Y1 }

// Detector runs the D-FINE text detection model.
type Detector struct {
	sess *ort.DynamicAdvancedSession
}

// NewDetector loads the text detector ONNX model.
func NewDetector(modelPath string) (*Detector, error) {
	sess, err := newSession(modelPath, []string{"images", "orig_target_sizes"}, []string{"labels", "boxes", "scores"})
	if err != nil {
		return nil, fmt.Errorf("load detector session: %w", err)
	}
	return &Detector{sess: sess}, nil
}

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

// Detect finds text bounding boxes in img with confidence >= confThreshold.
func (d *Detector) Detect(img image.Image, confThreshold float32) ([]TextBox, error) {
	bounds := img.Bounds()
	wOrig := bounds.Dx()
	hOrig := bounds.Dy()
	if wOrig <= 0 || hOrig <= 0 {
		return nil, nil
	}

	scale := math.Min(float64(detInputWidth)/float64(wOrig), float64(detInputHeight)/float64(hOrig))
	wResized := int(float64(wOrig) * scale)
	hResized := int(float64(hOrig) * scale)
	if wResized <= 0 || hResized <= 0 {
		return nil, nil
	}

	resized := image.NewRGBA(image.Rect(0, 0, wResized, hResized))
	draw.BiLinear.Scale(resized, resized.Bounds(), img, bounds, draw.Over, nil)

	// Prepare input tensor [1, 3, 544, 960] in BGR layout
	imgData := make([]float32, 1*3*detInputHeight*detInputWidth)
	strideH := detInputWidth
	strideC := detInputHeight * detInputWidth

	for y := 0; y < hResized; y++ {
		for x := 0; x < wResized; x++ {
			c := resized.RGBAAt(x, y)
			pxIdx := y*strideH + x
			// Channel 0: B, Channel 1: G, Channel 2: R
			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, detInputHeight, detInputWidth}, imgData)
	if err != nil {
		return nil, fmt.Errorf("create images tensor: %w", err)
	}
	defer imgTensor.Destroy()

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

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

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

	boxesData := boxesTensor.GetData()
	scoresData := scoresTensor.GetData()

	var boxes []TextBox
	for i, score := range scoresData {
		if score < confThreshold {
			continue
		}
		bOffset := i * 4
		x1 := int(math.Round(float64(boxesData[bOffset])))
		y1 := int(math.Round(float64(boxesData[bOffset+1])))
		x2 := int(math.Round(float64(boxesData[bOffset+2])))
		y2 := int(math.Round(float64(boxesData[bOffset+3])))

		// Clamp to image bounds
		x1 = clamp(x1, 0, wOrig)
		y1 = clamp(y1, 0, hOrig)
		x2 = clamp(x2, 0, wOrig)
		y2 = clamp(y2, 0, hOrig)

		w := x2 - x1
		h := y2 - y1
		if w <= 0 || h <= 0 {
			continue
		}

		boxes = append(boxes, TextBox{
			X1:         x1,
			Y1:         y1,
			X2:         x2,
			Y2:         y2,
			Score:      score,
			IsVertical: h > w,
		})
	}

	return boxes, nil
}

func clamp(val, low, high int) int {
	if val < low {
		return low
	}
	if val > high {
		return high
	}
	return val
}

// FilterRubyBoxes identifies and filters out small kana-only ruby/furigana boxes.
func FilterRubyBoxes(boxes []TextBox, isKanaOnly func(b TextBox) bool) []TextBox {
	if len(boxes) <= 1 {
		return boxes
	}

	drop := make(map[int]bool)

	for i, r := range boxes {
		for j, m := range boxes {
			if i == j {
				continue
			}

			// Case 1: Horizontal line with ruby above it
			if !m.IsVertical && !r.IsVertical {
				hM := m.Height()
				hR := r.Height()
				// Ruby height is significantly smaller
				if float64(hR) <= float64(hM)*0.7 {
					// Ruby sits immediately above main line
					if r.Y2 <= m.Y1+int(float64(hM)*0.35) && r.Y1 < m.Y1 && (m.Y1-r.Y2) <= int(float64(hM)*0.85) {
						oX := math.Max(0, math.Min(float64(r.X2), float64(m.X2))-math.Max(float64(r.X1), float64(m.X1)))
						if oX >= float64(r.Width())*0.55 {
							if isKanaOnly(r) {
								drop[i] = true
								break
							}
						}
					}
				}
			}

			// Case 2: Vertical column with ruby to the right
			if m.IsVertical && r.IsVertical {
				wM := m.Width()
				wR := r.Width()
				if float64(wR) <= float64(wM)*0.7 {
					// Ruby sits to the right of vertical column
					if r.X1 >= m.X1+int(float64(wM)*0.4) && (r.X1-m.X2) <= int(float64(wM)*0.85) {
						oY := math.Max(0, math.Min(float64(r.Y2), float64(m.Y2))-math.Max(float64(r.Y1), float64(m.Y1)))
						if oY >= float64(r.Height())*0.55 {
							if isKanaOnly(r) {
								drop[i] = true
								break
							}
						}
					}
				}
			}
		}
	}

	var filtered []TextBox
	for i, b := range boxes {
		if !drop[i] {
			filtered = append(filtered, b)
		}
	}
	return filtered
}