package ocr

import (
	"bytes"
	"fmt"
	"image"
	"image/png"

	"github.com/gen2brain/go-fitz"
	"github.com/otiai10/gosseract/v2"
)

type Box struct {
	Text string
	X    int
	Y    int
	W    int
	H    int
}

// RenderPage rasterizes the first page of a PDF.
func RenderPage(pdfData []byte, dpi int) (image.Image, int, int, error) {
	doc, err := fitz.NewFromMemory(pdfData)
	if err != nil {
		return nil, 0, 0, fmt.Errorf("open pdf: %w", err)
	}
	defer doc.Close()

	if doc.NumPage() == 0 {
		return nil, 0, 0, fmt.Errorf("pdf has no pages")
	}

	img, err := doc.ImageDPI(0, float64(dpi))
	if err != nil {
		return nil, 0, 0, fmt.Errorf("render page: %w", err)
	}
	bounds := img.Bounds()
	return img, bounds.Dx(), bounds.Dy(), nil
}

// OCRImage runs Tesseract word-level OCR on a rendered page image.
func (e *Engine) OCRImage(img image.Image) ([]Box, error) {
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		return nil, fmt.Errorf("encode png: %w", err)
	}

	client := e.Acquire()
	defer e.Release(client)

	if err := client.SetImageFromBytes(buf.Bytes()); err != nil {
		return nil, fmt.Errorf("tesseract image: %w", err)
	}

	boxes, err := client.GetBoundingBoxes(gosseract.RIL_WORD)
	if err != nil {
		return nil, fmt.Errorf("tesseract ocr: %w", err)
	}

	out := make([]Box, 0, len(boxes))
	for _, b := range boxes {
		text := trim(b.Word)
		if text == "" || b.Confidence <= 0 {
			continue
		}
		out = append(out, Box{
			Text: text,
			X:    b.Box.Min.X,
			Y:    b.Box.Min.Y,
			W:    b.Box.Dx(),
			H:    b.Box.Dy(),
		})
	}
	return out, nil
}

// ExtractBoxes renders the first PDF page and runs OCR with a pooled Tesseract client.
func (e *Engine) ExtractBoxes(pdfData []byte, dpi int) ([]Box, int, int, error) {
	img, pageW, pageH, err := RenderPage(pdfData, dpi)
	if err != nil {
		return nil, 0, 0, err
	}
	boxes, err := e.OCRImage(img)
	if err != nil {
		return nil, 0, 0, err
	}
	return boxes, pageW, pageH, nil
}

func trim(s string) string {
	start, end := 0, len(s)
	for start < end && (s[start] == ' ' || s[start] == '\t' || s[start] == '\n') {
		start++
	}
	for end > start && (s[end-1] == ' ' || s[end-1] == '\t' || s[end-1] == '\n') {
		end--
	}
	return s[start:end]
}