package table

import (
	"context"
	"image"
	pdf "ragflow/internal/deepdoc/parser/pdf/type"
	"testing"
)

// mockRotationDoc implements DocAnalyzer with deterministic OCR results per angle.
// The mock tracks the call sequence: EvaluateTableOrientation calls OCRRecognize
// once per angle in order 0°, 90°, 180°, 270°. Each call to OCRRecognize
// increments an internal counter and returns data for the corresponding angle.
type mockRotationDoc struct {
	// angle → {regions count, average confidence, error}
	angles map[int]struct {
		regions int
		avgConf float64
		err     error
	}
	callSeq int // incremented per OCRDetect call, selects the angle's data
}

var rotationOrder = []int{0, 90, 180, 270}

func (m *mockRotationDoc) DLA(_ context.Context, _ image.Image) ([]pdf.DLARegion, error) {
	return nil, nil
}
func (m *mockRotationDoc) TSR(_ context.Context, _ image.Image) ([]pdf.TSRCell, error) {
	return nil, nil
}
func (m *mockRotationDoc) OCR(_ image.Image) (string, error) { return "", nil }
func (m *mockRotationDoc) Health() bool                      { return true }

func (m *mockRotationDoc) OCRDetect(_ context.Context, _ image.Image) ([]pdf.OCRBox, error) {
	// EvaluateTableOrientation scores by OCRRecognize; detection output is
	// unused here. Return empty to satisfy the DocAnalyzer interface.
	return nil, nil
}

func (m *mockRotationDoc) OCRRecognize(_ context.Context, _ image.Image) ([]pdf.OCRText, error) {
	// EvaluateTableOrientation calls OCRRecognize once per angle in order
	// 0°, 90°, 180°, 270°. Track the call sequence here so each call returns
	// the recognition result for the corresponding angle.
	angle := rotationOrder[m.callSeq%len(rotationOrder)]
	m.callSeq++
	cfg, ok := m.angles[angle]
	if !ok {
		cfg = m.angles[0]
	}
	if cfg.err != nil {
		return nil, cfg.err
	}
	if cfg.regions == 0 {
		return nil, nil
	}
	texts := make([]pdf.OCRText, cfg.regions)
	for i := 0; i < cfg.regions; i++ {
		texts[i] = pdf.OCRText{Text: "X", Confidence: cfg.avgConf}
	}
	return texts, nil
}

func makeTestTableImage() image.Image {
	return image.NewRGBA(image.Rect(0, 0, 200, 100))
}

func TestEvaluateTableOrientation(t *testing.T) {
	t.Run("normal table 0° wins", func(t *testing.T) {
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0: {regions: 10, avgConf: 0.9},
			},
		}
		angle, _, scores := EvaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 0 {
			t.Errorf("expected 0°, got %d° (scores: %v)", angle, scores)
		}
	})

	t.Run("90° rotated table wins", func(t *testing.T) {
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:   {regions: 2, avgConf: 0.2},
				90:  {regions: 10, avgConf: 0.9},
				180: {regions: 2, avgConf: 0.2},
				270: {regions: 2, avgConf: 0.2},
			},
		}
		angle, _, scores := EvaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 90 {
			t.Errorf("expected 90°, got %d° (scores: %v)", angle, scores)
		}
	})

	t.Run("180° rotated table wins", func(t *testing.T) {
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:   {regions: 1, avgConf: 0.1},
				90:  {regions: 1, avgConf: 0.1},
				180: {regions: 8, avgConf: 0.85},
				270: {regions: 1, avgConf: 0.1},
			},
		}
		angle, _, scores := EvaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 180 {
			t.Errorf("expected 180°, got %d° (scores: %v)", angle, scores)
		}
	})

	t.Run("270° rotated table wins", func(t *testing.T) {
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:   {regions: 1, avgConf: 0.1},
				90:  {regions: 1, avgConf: 0.1},
				180: {regions: 1, avgConf: 0.1},
				270: {regions: 9, avgConf: 0.88},
			},
		}
		angle, _, scores := EvaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 270 {
			t.Errorf("expected 270°, got %d° (scores: %v)", angle, scores)
		}
	})

	t.Run("threshold protection — 0° keeps when confidence diff too small", func(t *testing.T) {
		// Recognition scores 0.50 vs 0.55 are too close (< 0.2 margin) → 0° wins.
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:  {regions: 8, avgConf: 0.50},
				90: {regions: 8, avgConf: 0.55},
			},
		}
		angle, _, _ := EvaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 0 {
			t.Errorf("expected 0° (threshold protection), got %d°", angle)
		}
	})

	t.Run("threshold pass — 90° wins when recognition confidence is clearly higher", func(t *testing.T) {
		// 0° reads poorly (0.30) AND 90° reads well (0.90) → 90° wins.
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:  {regions: 4, avgConf: 0.30},
				90: {regions: 10, avgConf: 0.90},
			},
		}
		angle, _, _ := EvaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 90 {
			t.Errorf("expected 90° (threshold passed), got %d°", angle)
		}
	})

	t.Run("threshold guard — score_0 >= 0.8 blocks rotation despite large margin", func(t *testing.T) {
		// Isolate the score_0 < 0.8 clause: 0° reads well (0.80) and 90° is
		// clearly higher (1.00), so the margin clause (combined diff 0.22 > 0.2)
		// passes, but score_0 = 0.88 >= 0.8 must still force keeping 0°.
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:  {regions: 50, avgConf: 0.80},
				90: {regions: 50, avgConf: 1.00},
			},
		}
		angle, _, _ := EvaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 0 {
			t.Errorf("expected 0° (score_0 >= 0.8 guard), got %d°", angle)
		}
	})

	t.Run("all angles fail OCR → fallback 0°", func(t *testing.T) {
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:   {err: errMockOCR},
				90:  {err: errMockOCR},
				180: {err: errMockOCR},
				270: {err: errMockOCR},
			},
		}
		angle, img, scores := EvaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 0 {
			t.Errorf("expected 0° fallback, got %d°", angle)
		}
		if img == nil {
			t.Error("expected non-nil fallback image")
		}
		for _, s := range scores {
			if s != 0 {
				t.Error("all scores should be 0 on OCR failure")
			}
		}
	})

	t.Run("zero score_0 with low non-zero score — keep 0°", func(t *testing.T) {
		// 0° has no recognized text (score_0 == 0). A non-zero angle with a
		// low combined score must NOT be accepted, matching Python's
		// `score_0 is not None` threshold (not `score_0 > 0`).
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:  {regions: 0, avgConf: 0},
				90: {regions: 2, avgConf: 0.05},
			},
		}
		angle, _, _ := EvaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 0 {
			t.Errorf("expected 0° (score_0 == 0, low non-zero score), got %d°", angle)
		}
	})

	t.Run("zero score_0 with high non-zero score — accept rotation", func(t *testing.T) {
		// 0° has no recognized text (score_0 == 0) but 90° reads clearly
		// (combined 1.045 > 0.2). Mirrors Python: score_0 is not None, so the
		// margin clause alone decides and 90° is accepted.
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:  {regions: 0, avgConf: 0},
				90: {regions: 50, avgConf: 0.95},
			},
		}
		angle, _, _ := EvaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 90 {
			t.Errorf("expected 90° (score_0 == 0, high non-zero score), got %d°", angle)
		}
	})
}

var errMockOCR = &mockError{"mock OCR failure"}

type mockError struct{ msg string }

func (e *mockError) Error() string { return e.msg }