package onnx

// Pipeline combines meikiocr's text detector (D-FINE) with meikiocr's horizontal
// recognizer and manga-ocr's vertical reader into a single unified OCR engine.
//
// Attributions:
// - meikiocr: https://github.com/rtr46/meikiocr (Apache-2.0)
// - meikiocr weights: rtr46/meiki.text.detect.v0 & meiki.txt.recognition.v0 (LGPL-3.0)
// - manga-ocr: https://github.com/kha-white/manga-ocr (Apache-2.0)
// - manga-ocr ONNX: mayocream/manga-ocr-onnx (Apache-2.0)

import (
	"fmt"
	"image"
	"path/filepath"
	"sort"

	"chirami/internal/ocr"
	"golang.org/x/image/draw"
)

type Pipeline struct {
	det   *Detector
	meiki *MeikiRecognizer
	manga *MangaOCR
}

// NewPipeline initializes all models from modelDir.
func NewPipeline(modelDir string) (*Pipeline, error) {
	detPath := filepath.Join(modelDir, "meiki.text.detect.v0.1.960x544.onnx")
	recPath := filepath.Join(modelDir, "meiki.text.rec.v0.960x32.onnx")
	encPath := filepath.Join(modelDir, "encoder_model.onnx")
	decPath := filepath.Join(modelDir, "decoder_model.onnx")
	vocabPath := filepath.Join(modelDir, "vocab.txt")

	det, err := NewDetector(detPath)
	if err != nil {
		return nil, fmt.Errorf("detector: %w", err)
	}

	meiki, err := NewMeikiRecognizer(recPath)
	if err != nil {
		det.Close()
		return nil, fmt.Errorf("meiki recognizer: %w", err)
	}

	manga, err := NewMangaOCR(encPath, decPath, vocabPath)
	if err != nil {
		det.Close()
		meiki.Close()
		return nil, fmt.Errorf("manga-ocr: %w", err)
	}

	return &Pipeline{
		det:   det,
		meiki: meiki,
		manga: manga,
	}, nil
}

// Close releases all underlying ONNX sessions.
func (p *Pipeline) Close() error {
	var errs []error
	if p.det != nil {
		if err := p.det.Close(); err != nil {
			errs = append(errs, err)
		}
		p.det = nil
	}
	if p.meiki != nil {
		if err := p.meiki.Close(); err != nil {
			errs = append(errs, err)
		}
		p.meiki = nil
	}
	if p.manga != nil {
		if err := p.manga.Close(); err != nil {
			errs = append(errs, err)
		}
		p.manga = nil
	}
	if len(errs) > 0 {
		return fmt.Errorf("close pipeline: %v", errs)
	}
	return nil
}

// Recognize processes an image and returns recognized text lines in reading order.
func (p *Pipeline) Recognize(img image.Image) (ocr.Result, error) {
	boxes, err := p.det.Detect(img, defaultDetConf)
	if err != nil {
		return ocr.Result{Engine: "onnx", Error: err.Error()}, err
	}

	// If detector found nothing, fallback to manga-ocr on the whole image
	if len(boxes) == 0 {
		text, err := p.manga.Recognize(img)
		if err != nil {
			return ocr.Result{Engine: "onnx", Error: err.Error()}, err
		}
		if text == "" {
			return ocr.Result{Engine: "onnx"}, nil
		}
		return ocr.Result{
			Lines:  []ocr.Line{{Text: text, Vertical: false}},
			Engine: "onnx",
		}, nil
	}

	// Filter furigana/ruby boxes so they don't end up inside the output text
	filteredBoxes := FilterRubyBoxes(boxes, func(b TextBox) bool {
		crop := cropImage(img, b.X1, b.Y1, b.X2, b.Y2)
		var txt string
		if b.IsVertical {
			txt, _ = p.manga.Recognize(crop)
		} else {
			txt, _ = p.meiki.RecognizeLine(crop)
		}
		return IsKanaOnly(txt)
	})

	// Separate into horizontal and vertical boxes
	var hBoxes []TextBox
	var vBoxes []TextBox
	for _, b := range filteredBoxes {
		if b.IsVertical {
			vBoxes = append(vBoxes, b)
		} else {
			hBoxes = append(hBoxes, b)
		}
	}

	// Reading order:
	// Horizontal: top-to-bottom (ascending Y1)
	sort.Slice(hBoxes, func(i, j int) bool {
		return hBoxes[i].Y1 < hBoxes[j].Y1
	})

	// Vertical: right-to-left (descending X1)
	sort.Slice(vBoxes, func(i, j int) bool {
		return vBoxes[i].X1 > vBoxes[j].X1
	})

	var lines []ocr.Line

	// 1. Process horizontal lines with meiki recognizer
	for _, b := range hBoxes {
		crop := cropImage(img, b.X1, b.Y1, b.X2, b.Y2)
		txt, err := p.meiki.RecognizeLine(crop)
		if err != nil {
			return ocr.Result{Engine: "onnx", Error: err.Error()}, err
		}
		if txt != "" {
			lines = append(lines, ocr.Line{Text: txt, Vertical: false})
		}
	}

	// 2. Process vertical columns with manga-ocr
	for _, b := range vBoxes {
		crop := cropImage(img, b.X1, b.Y1, b.X2, b.Y2)
		txt, err := p.manga.Recognize(crop)
		if err != nil {
			return ocr.Result{Engine: "onnx", Error: err.Error()}, err
		}
		if txt != "" {
			lines = append(lines, ocr.Line{Text: txt, Vertical: true})
		}
	}

	return ocr.Result{
		Lines:  lines,
		Engine: "onnx",
	}, nil
}

func cropImage(img image.Image, x1, y1, x2, y2 int) image.Image {
	type subImager interface {
		SubImage(r image.Rectangle) image.Image
	}
	rect := image.Rect(x1, y1, x2, y2)
	if si, ok := img.(subImager); ok {
		return si.SubImage(rect)
	}
	dst := image.NewRGBA(image.Rect(0, 0, x2-x1, y2-y1))
	draw.Draw(dst, dst.Bounds(), img, rect.Min, draw.Src)
	return dst
}