package onnx

// Ported from manga-ocr (https://github.com/kha-white/manga-ocr), Apache-2.0.
// Model weights: mayocream/manga-ocr-onnx (kha-white base ONNX), Apache-2.0.

import (
	"bufio"
	"fmt"
	"image"
	"os"
	"regexp"
	"strings"

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

const (
	mangaInputSize = 224
	maxDecodeLen   = 100
	bosTokenID     = 2 // [CLS]
	eosTokenID     = 3 // [SEP]
)

var dotRegex = regexp.MustCompile(`[・.]{2,}`)

// MangaOCR runs the VisionEncoderDecoder manga-ocr model.
type MangaOCR struct {
	sessEnc *ort.DynamicAdvancedSession
	sessDec *ort.DynamicAdvancedSession
	vocab   []string
}

// NewMangaOCR loads the encoder, decoder, and vocabulary.
func NewMangaOCR(encPath, decPath, vocabPath string) (*MangaOCR, error) {
	vocab, err := loadVocab(vocabPath)
	if err != nil {
		return nil, fmt.Errorf("load vocab: %w", err)
	}

	sessEnc, err := newSession(encPath, []string{"pixel_values"}, []string{"last_hidden_state"})
	if err != nil {
		return nil, fmt.Errorf("load manga encoder: %w", err)
	}

	sessDec, err := newSession(decPath, []string{"input_ids", "encoder_hidden_states"}, []string{"logits"})
	if err != nil {
		sessEnc.Destroy()
		return nil, fmt.Errorf("load manga decoder: %w", err)
	}

	return &MangaOCR{
		sessEnc: sessEnc,
		sessDec: sessDec,
		vocab:   vocab,
	}, nil
}

func loadVocab(path string) ([]string, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var vocab []string
	sc := bufio.NewScanner(f)
	for sc.Scan() {
		vocab = append(vocab, sc.Text())
	}
	if err := sc.Err(); err != nil {
		return nil, err
	}
	return vocab, nil
}

// Close frees all session resources.
func (m *MangaOCR) Close() error {
	var errs []error
	if m.sessEnc != nil {
		if err := m.sessEnc.Destroy(); err != nil {
			errs = append(errs, err)
		}
		m.sessEnc = nil
	}
	if m.sessDec != nil {
		if err := m.sessDec.Destroy(); err != nil {
			errs = append(errs, err)
		}
		m.sessDec = nil
	}
	if len(errs) > 0 {
		return fmt.Errorf("close manga-ocr: %v", errs)
	}
	return nil
}

// Recognize processes a single image crop or bubble.
func (m *MangaOCR) Recognize(img image.Image) (string, error) {
	bounds := img.Bounds()
	if bounds.Dx() <= 0 || bounds.Dy() <= 0 {
		return "", nil
	}

	// 1. Grayscale and resize to 224x224 bilinear
	resized := image.NewRGBA(image.Rect(0, 0, mangaInputSize, mangaInputSize))
	draw.BiLinear.Scale(resized, resized.Bounds(), img, bounds, draw.Over, nil)

	// 2. Normalize: (gray / 255.0 - 0.5) / 0.5 replicated across 3 RGB channels
	pixData := make([]float32, 1*3*mangaInputSize*mangaInputSize)
	channelStride := mangaInputSize * mangaInputSize

	for y := 0; y < mangaInputSize; y++ {
		for x := 0; x < mangaInputSize; x++ {
			c := resized.RGBAAt(x, y)
			// Convert RGB to grayscale (ITU-R 601-2 luma)
			gray := 0.299*float32(c.R) + 0.587*float32(c.G) + 0.114*float32(c.B)
			norm := (gray/255.0 - 0.5) / 0.5
			idx := y*mangaInputSize + x
			pixData[0*channelStride+idx] = norm
			pixData[1*channelStride+idx] = norm
			pixData[2*channelStride+idx] = norm
		}
	}

	pixTensor, err := ort.NewTensor(ort.Shape{1, 3, mangaInputSize, mangaInputSize}, pixData)
	if err != nil {
		return "", fmt.Errorf("create pixel_values tensor: %w", err)
	}
	defer pixTensor.Destroy()

	// 3. Encoder inference
	encOut := make([]ort.Value, 1)
	err = m.sessEnc.Run([]ort.Value{pixTensor}, encOut)
	if err != nil {
		return "", fmt.Errorf("encoder run: %w", err)
	}
	defer encOut[0].Destroy()

	hiddenState := encOut[0].(*ort.Tensor[float32])

	// 4. Greedy autoregressive decoding
	inputIDs := []int64{bosTokenID}
	for step := 0; step < maxDecodeLen; step++ {
		seqLen := int64(len(inputIDs))
		idsTensor, err := ort.NewTensor(ort.Shape{1, seqLen}, inputIDs)
		if err != nil {
			return "", fmt.Errorf("create input_ids tensor: %w", err)
		}

		decOut := make([]ort.Value, 1)
		err = m.sessDec.Run([]ort.Value{idsTensor, hiddenState}, decOut)
		idsTensor.Destroy()
		if err != nil {
			return "", fmt.Errorf("decoder run step %d: %w", step, err)
		}

		logitsTensor := decOut[0].(*ort.Tensor[float32])
		logits := logitsTensor.GetData()
		vocabSize := int(logitsTensor.GetShape()[2])
		lastLogits := logits[(seqLen-1)*int64(vocabSize) : seqLen*int64(vocabSize)]

		// Argmax
		var bestTok int64 = 0
		var bestVal float32 = -1e9
		for tok, val := range lastLogits {
			if val > bestVal {
				bestVal = val
				bestTok = int64(tok)
			}
		}
		decOut[0].Destroy()

		inputIDs = append(inputIDs, bestTok)
		if bestTok == eosTokenID {
			break
		}
	}

	// 5. Detokenize and postprocess
	var tokens []string
	for _, id := range inputIDs {
		if id >= 0 && id <= 4 { // [PAD], [UNK], [CLS], [SEP], [MASK]
			continue
		}
		if int(id) < len(m.vocab) {
			tokens = append(tokens, m.vocab[id])
		}
	}

	text := strings.Join(tokens, "")
	// Strip spaces between tokens
	text = strings.Join(strings.Fields(text), "")
	// Standardize ellipsis and dots
	text = strings.ReplaceAll(text, "…", "...")
	text = dotRegex.ReplaceAllStringFunc(text, func(m string) string {
		return strings.Repeat(".", len([]rune(m)))
	})
	// Full-width conversion
	text = H2Z(text)

	return text, nil
}