package documenttext

import (
	"context"
	"encoding/json"
	"io"
	"math"
	"os"
	"strings"
	"testing"

	documenttypes "queryorchestration/internal/document/types"
	textractmock "queryorchestration/mocks/textract"

	"github.com/aws/aws-sdk-go-v2/service/textract"
	"github.com/aws/aws-sdk-go-v2/service/textract/types"
	"github.com/google/go-cmp/cmp"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

func TestGetBasePage(t *testing.T) {
	ctx := context.Background()
	cfg := DocTextConfig{}
	svc := Service{
		cfg: &cfg,
	}

	mockTextract := textractmock.NewMockTextractClient(t)
	cfg.TextractClient = mockTextract
	sent := 0

	pdf, err := documenttypes.NewPDFFromReader(strings.NewReader(pdfHelloWorld))
	require.NoError(t, err)

	textractFile := "../../../assets/sampleGeneration/generated/helloWorld.gen"
	textractBaseOut := getTextractFileResponse(t, textractFile)

	mockTextract.EXPECT().
		AnalyzeDocument(
			mock.Anything,
			mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
				if sent == 0 {
					sent++
					return true
				}
				return false
			}),
			mock.Anything,
		).
		Return(&textractBaseOut["base"][0], nil).Once()

	text, err := svc.GetBasePage(ctx, pdf, 0)
	require.NoError(t, err)
	assert.EqualExportedValues(t, textractBaseOut["base"][0], text)
}

func TestGetDocumentText(t *testing.T) {
	ctx := context.Background()
	cfg := DocTextConfig{}
	svc := Service{
		cfg: &cfg,
	}

	t.Run("Hello World", func(t *testing.T) {
		mockTextract := textractmock.NewMockTextractClient(t)
		cfg.TextractClient = mockTextract
		sent := 0

		pdf, err := documenttypes.NewPDFFromReader(strings.NewReader(pdfHelloWorld))
		require.NoError(t, err)

		textractFile := "../../../assets/sampleGeneration/generated/helloWorld.gen"
		textractBaseOut := getTextractFileResponse(t, textractFile)

		mockTextract.EXPECT().
			AnalyzeDocument(
				mock.Anything,
				mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
					if sent == 0 {
						sent++
						return true
					}
					return false
				}),
				mock.Anything,
			).
			Return(&textractBaseOut["base"][0], nil).Once()

		text, err := svc.getDocumentText(ctx, pdf)
		require.NoError(t, err)
		expected := `Start of Page No. = 1

Hello World!


`
		assertStringToReader(t, expected, text)
	})

	t.Run("sampleOne - cnc_01-06", func(t *testing.T) {
		mockTextract := textractmock.NewMockTextractClient(t)
		cfg.TextractClient = mockTextract
		sent := 0

		textractFile := "../../../assets/sampleOne/generated/cnc_01-06.gen"
		textractBaseOut := getTextractFileResponse(t, textractFile)

		pdfFile, err := os.Open("../../../assets/sampleOne/cnc_01-06.pdf")
		require.NoError(t, err)
		defer pdfFile.Close()

		pdf, err := documenttypes.NewPDFFromReader(pdfFile)
		require.NoError(t, err)

		txtFile, err := os.Open("../../../assets/sampleOne/cnc_01-06.txt")
		require.NoError(t, err)
		defer txtFile.Close()

		mockTextract.EXPECT().
			AnalyzeDocument(
				mock.Anything,
				mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
					if sent == 0 {
						sent++
						return true
					}
					return false
				}),
				mock.Anything,
			).
			Return(&textractBaseOut["base"][0], nil).Once()

		mockTextract.EXPECT().
			AnalyzeDocument(
				mock.Anything,
				mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
					if sent == 1 {
						sent++
						return true
					}
					return false
				}),
				mock.Anything,
			).
			Return(&textractBaseOut["base"][1], nil).Once()

		mockTextract.EXPECT().
			AnalyzeDocument(
				mock.Anything,
				mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
					if sent == 2 {
						sent++
						return true
					}
					return false
				}),
				mock.Anything,
			).
			Return(&textractBaseOut["base"][2], nil).Once()

		mockTextract.EXPECT().
			AnalyzeDocument(
				mock.Anything,
				mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
					if sent == 3 {
						sent++
						return true
					}
					return false
				}),
				mock.Anything,
			).
			Return(&textractBaseOut["base"][3], nil).Once()

		text, err := svc.getDocumentText(ctx, pdf)
		require.NoError(t, err)

		assertReaders(t, txtFile, text)
	})
	t.Run("sampleOne - hn_0109", func(t *testing.T) {
		mockTextract := textractmock.NewMockTextractClient(t)
		cfg.TextractClient = mockTextract
		sent := 0

		textractFile := "../../../assets/sampleOne/generated/hn_0109.gen"
		textractBaseOut := getTextractFileResponse(t, textractFile)

		pdfFile, err := os.Open("../../../assets/sampleOne/hn_0109.pdf")
		require.NoError(t, err)
		defer pdfFile.Close()

		pdf, err := documenttypes.NewPDFFromReader(pdfFile)
		require.NoError(t, err)

		txtFile, err := os.Open("../../../assets/sampleOne/hn_0109.txt")
		require.NoError(t, err)
		defer txtFile.Close()

		mockTextract.EXPECT().
			AnalyzeDocument(
				mock.Anything,
				mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
					if sent == 0 {
						sent++
						return true
					}
					return false
				}),
				mock.Anything,
			).
			Return(&textractBaseOut["base"][0], nil).Once()

		mockTextract.EXPECT().
			AnalyzeDocument(
				mock.Anything,
				mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
					if sent == 1 {
						sent++
						return true
					}
					return false
				}),
				mock.Anything,
			).
			Return(&textractBaseOut["base"][1], nil).Once()

		mockTextract.EXPECT().
			AnalyzeDocument(
				mock.Anything,
				mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
					if sent == 2 && in.FeatureTypes[0] == types.FeatureTypeTables {
						sent++
						return true
					}
					return false
				}),
				mock.Anything,
			).
			Return(&textractBaseOut["table"][0], nil).Once()

		mockTextract.EXPECT().
			AnalyzeDocument(
				mock.Anything,
				mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
					if sent == 3 {
						sent++
						return true
					}
					return false
				}),
				mock.Anything,
			).
			Return(&textractBaseOut["base"][2], nil).Once()

		mockTextract.EXPECT().
			AnalyzeDocument(
				mock.Anything,
				mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
					if sent == 4 && in.FeatureTypes[0] == types.FeatureTypeTables {
						sent++
						return true
					}
					return false
				}),
				mock.Anything,
			).
			Return(&textractBaseOut["table"][1], nil).Once()

		mockTextract.EXPECT().
			AnalyzeDocument(
				mock.Anything,
				mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
					if sent == 5 {
						sent++
						return true
					}
					return false
				}),
				mock.Anything,
			).
			Return(&textractBaseOut["base"][3], nil).Once()

		text, err := svc.getDocumentText(ctx, pdf)
		require.NoError(t, err)

		assertReaders(t, txtFile, text)
	})
}

func TestGetBlockMap(t *testing.T) {
	svc := Service{}
	t.Run("0 blocks", func(t *testing.T) {
		blocks := []types.Block{}
		outBlocks := svc.getBlockMap(blocks)
		assert.Equal(t, map[uuid.UUID]types.Block{}, outBlocks)
	})
	t.Run("nil id blocks", func(t *testing.T) {
		blocks := []types.Block{
			{},
		}
		outBlocks := svc.getBlockMap(blocks)
		assert.Equal(t, map[uuid.UUID]types.Block{}, outBlocks)
	})
	t.Run("1 block", func(t *testing.T) {
		blockId := uuid.New()
		blockIdStr := blockId.String()
		blocks := []types.Block{
			{
				Id: &blockIdStr,
			},
		}
		outBlocks := svc.getBlockMap(blocks)
		assert.Equal(t, map[uuid.UUID]types.Block{
			blockId: {
				Id: &blockIdStr,
			},
		}, outBlocks)
	})
	t.Run("invalid block id", func(t *testing.T) {
		blockIdStr := "blockId.String()"
		blocks := []types.Block{
			{
				Id: &blockIdStr,
			},
		}
		outBlocks := svc.getBlockMap(blocks)
		assert.Equal(t, map[uuid.UUID]types.Block{}, outBlocks)
	})
	t.Run("2 blocks", func(t *testing.T) {
		blockOneId := uuid.New()
		blockOneIdStr := blockOneId.String()
		blockTwoId := uuid.New()
		blockTwoIdStr := blockTwoId.String()
		blocks := []types.Block{
			{
				Id: &blockOneIdStr,
			},
			{
				Id: &blockTwoIdStr,
			},
		}
		outBlocks := svc.getBlockMap(blocks)
		assert.Equal(t, map[uuid.UUID]types.Block{
			blockOneId: {
				Id: &blockOneIdStr,
			},
			blockTwoId: {
				Id: &blockTwoIdStr,
			},
		}, outBlocks)
	})
}

func TestMergeDocument(t *testing.T) {
	svc := Service{}
	t.Run("0 pages", func(t *testing.T) {
		pages := []string{}
		out, err := svc.mergeDocument(pages)
		require.NoError(t, err)
		outStr, err := io.ReadAll(out)
		require.NoError(t, err)
		assert.Equal(t, ``, string(outStr))
	})
	t.Run("1 page", func(t *testing.T) {
		pages := []string{
			"page one",
		}
		out, err := svc.mergeDocument(pages)
		require.NoError(t, err)
		outStr, err := io.ReadAll(out)
		require.NoError(t, err)
		assert.Equal(t, `Start of Page No. = 1

page one

`, string(outStr))
	})
	t.Run("2 pages", func(t *testing.T) {
		pages := []string{
			"page one",
			"page two",
		}
		out, err := svc.mergeDocument(pages)
		require.NoError(t, err)
		outStr, err := io.ReadAll(out)
		require.NoError(t, err)
		assert.Equal(t, `Start of Page No. = 1

page one

Start of Page No. = 2

page two

`, string(outStr))
	})
	t.Run("3 pages", func(t *testing.T) {
		pages := []string{
			"page one",
			"page two",
			"page three",
		}
		out, err := svc.mergeDocument(pages)
		require.NoError(t, err)
		outStr, err := io.ReadAll(out)
		require.NoError(t, err)
		assert.Equal(t, `Start of Page No. = 1

page one

Start of Page No. = 2

page two

Start of Page No. = 3

page three

`, string(outStr))
	})
}

func getTextractFileResponse(t *testing.T, filename string) map[string][]textract.AnalyzeDocumentOutput {
	file, err := os.Open(filename)
	require.NoError(t, err)
	defer file.Close()

	decoder := json.NewDecoder(file)
	var out map[string][]textract.AnalyzeDocumentOutput
	err = decoder.Decode(&out)
	require.NoError(t, err)

	return out
}

func assertStringToReader(t *testing.T, expected string, actual io.Reader) {
	actualBytes, err := io.ReadAll(actual)
	require.NoError(t, err)
	actualStr := string(actualBytes)
	if !assert.Equal(t, expected, actualStr) {
		diff := cmp.Diff(expected, actualStr)
		t.Logf("Detailed diff (-expected +actual):\n%s", diff)
	}
}

func assertReaders(t *testing.T, expected io.Reader, actual io.Reader) {
	actualBytes, err := io.ReadAll(actual)
	require.NoError(t, err)
	actualStr := string(actualBytes)

	expectedBytes, err := io.ReadAll(expected)
	require.NoError(t, err)
	expectedStr := string(expectedBytes)

	if !assert.Equal(t, expectedStr, actualStr) {
		diff := cmp.Diff(expectedStr, actualStr)
		t.Logf("Detailed diff (-expected +actual):\n%s", diff)
	}
}

func levenshteinDistance(s1, s2 string) int {
	rows, cols := len(s1)+1, len(s2)+1
	dist := make([][]int, rows)
	for i := range dist {
		dist[i] = make([]int, cols)
	}

	for i := 0; i < rows; i++ {
		dist[i][0] = i
	}
	for j := 0; j < cols; j++ {
		dist[0][j] = j
	}

	for i := 1; i < rows; i++ {
		for j := 1; j < cols; j++ {
			cost := 1
			if s1[i-1] == s2[j-1] {
				cost = 0
			}
			dist[i][j] = min(
				dist[i-1][j]+1,
				dist[i][j-1]+1,
				dist[i-1][j-1]+cost,
			)
		}
	}

	return dist[rows-1][cols-1]
}

func calculateSimilarityPercentage(s1, s2 string) float64 {
	distance := levenshteinDistance(s1, s2)
	maxLen := math.Max(float64(len(s1)), float64(len(s2)))

	if maxLen == 0 {
		return 100.0
	}

	return 100.0 * (1.0 - float64(distance)/maxLen)
}

func AreSimilar(s1, s2 string, threshold float64) bool {
	return calculateSimilarityPercentage(s1, s2) >= threshold
}

const pdfHelloWorld = `%PDF-1.4
%����
1 0 obj
<<
  /Type /Catalog
  /Pages 2 0 R
  /Version /1.4
>>
endobj
2 0 obj
<<
  /Type /Pages
  /Kids [3 0 R]
  /Count 1
>>
endobj
3 0 obj
<<
  /Type /Page
  /Parent 2 0 R
  /MediaBox [0 0 612 792]
  /Resources <<
    /Font <<
      /F1 <<
        /Type /Font
        /Subtype /Type1
        /BaseFont /Helvetica
      >>
    >>
  >>
  /Contents 4 0 R
>>
endobj
4 0 obj
<<
  /Length
  44
>>
stream
BT
/F1 24 Tf
100 700 Td
(Hello World!) Tj
ET
endstream
endobj
xref
0 5
0000000000 65535 f
0000000015 00000 n
0000000086 00000 n
0000000151 00000 n
0000000376 00000 n
trailer
<<
  /Size 5
  /Root 1 0 R
>>
startxref
472
%%EOF`