query-orchestration/api/queryAPI/documents_test.go

package queryapi_test

import (
	"archive/zip"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/google/uuid"
	"github.com/labstack/echo/v4"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	queryapi "queryorchestration/api/queryAPI"
	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/test"
)

func TestUploadDocument(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)
	acfg := test.CreateAWSContainer(t, cfg)

	test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
	test.CreateBucket(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: "client_id",
		Name:     "client_name",
	})
	require.NoError(t, err)

	var buf bytes.Buffer
	writer := multipart.NewWriter(&buf)

	fileWriter, err := writer.CreateFormFile("file", filepath.Base("../../assets/original/helloWorld.pdf"))
	require.NoError(t, err)

	_, err = fileWriter.Write([]byte("hi"))
	require.NoError(t, err)

	err = writer.Close()
	require.NoError(t, err)

	ctx, rec := createContextWithBody(t, buf.Bytes())

	ctx.Request().Header.Set("Content-Type", writer.FormDataContentType())

	err = cons.UploadDocument(ctx, "client_id")
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)
}

func TestListDocumentsByClientId(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: "client_id",
		Name:     "client_name",
	})
	require.NoError(t, err)
	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
		Clientid: "client_id",
		Hash:     "hash",
		BatchID:  nil,
	})
	require.NoError(t, err)

	ctx, rec := createContext(t)

	err = cons.ListDocumentsByClientId(ctx, "client_id")
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)
	assertBody(t, rec, queryapi.ListDocuments{
		queryapi.DocumentSummary{
			Id:   docId,
			Hash: "hash",
		},
	})
}

func TestGetDocument(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: "client_id",
		Name:     "client_name",
	})
	require.NoError(t, err)
	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
		Clientid: "client_id",
		Hash:     "hash",
		BatchID:  nil,
	})
	require.NoError(t, err)

	ctx, rec := createContext(t)

	err = cons.GetDocument(ctx, docId)
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)
	assertBody(t, rec, queryapi.Document{
		Id:       docId,
		ClientId: "client_id",
		Hash:     "hash",
		Fields:   map[string]interface{}{},
	})
}

func TestListDocumentBatches(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	// Create test client
	clientID := "test_client_list"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Batch List",
	})
	require.NoError(t, err)

	// Create test batch
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test listing batches with default parameters
	ctx, rec := createContext(t)
	err = cons.ListDocumentBatches(ctx, clientID, queryapi.ListDocumentBatchesParams{})
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)

	response := getBody[queryapi.BatchUploadList](t, rec)
	assert.Len(t, response.Batches, 1)
	assert.Equal(t, batchID.String(), response.Batches[0].BatchId.String())

	// Test listing batches with custom limit and offset
	limit := int32(5)
	offset := int32(0)
	ctx2, rec2 := createContext(t)
	err = cons.ListDocumentBatches(ctx2, clientID, queryapi.ListDocumentBatchesParams{
		Limit:  &limit,
		Offset: &offset,
	})
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec2.Code)

	response2 := getBody[queryapi.BatchUploadList](t, rec2)
	assert.Len(t, response2.Batches, 1)
	assert.Equal(t, batchID.String(), response2.Batches[0].BatchId.String())
}

func TestUploadDocumentBatch(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)
	test.CreateAWSResources(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	// Create test client
	clientID := "test_client_upload"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Batch Upload",
	})
	require.NoError(t, err)

	// Create multipart form with ZIP file
	var buf bytes.Buffer
	writer := multipart.NewWriter(&buf)
	fileWriter, err := writer.CreateFormFile("archive", "test.zip")
	require.NoError(t, err)
	_, err = fileWriter.Write([]byte("PK")) // Minimal ZIP file header
	require.NoError(t, err)
	err = writer.Close()
	require.NoError(t, err)

	ctx, rec := createContextWithBody(t, buf.Bytes())
	ctx.Request().Header.Set("Content-Type", writer.FormDataContentType())

	// Test batch upload
	err = cons.UploadDocumentBatch(ctx, clientID)
	require.NoError(t, err)
	assert.Equal(t, http.StatusAccepted, rec.Code)

	response := getBody[queryapi.BatchUploadResponse](t, rec)
	assert.NotEmpty(t, response.BatchId)
	assert.Equal(t, queryapi.BatchStatusProcessing, response.Status)
}

func TestGetDocumentBatch(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	// Create test client
	clientID := "test_client_get"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Batch Get",
	})
	require.NoError(t, err)

	// Create test batch
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test getting batch details
	ctx, rec := createContext(t)
	err = cons.GetDocumentBatch(ctx, clientID, queryapi.BatchID(batchID))
	require.NoError(t, err)
	assert.Equal(t, http.StatusOK, rec.Code)

	response := getBody[queryapi.BatchUploadDetails](t, rec)
	assert.Equal(t, batchID.String(), response.BatchId.String())
	assert.Equal(t, "test.zip", response.OriginalFilename)
}

func TestCancelDocumentBatch(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	// Create test client
	clientID := "test_client_cancel"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Batch Cancel",
	})
	require.NoError(t, err)

	// Create test batch
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test canceling batch
	ctx, rec := createContext(t)
	err = cons.CancelDocumentBatch(ctx, clientID, queryapi.BatchID(batchID))
	require.NoError(t, err)
	assert.Equal(t, http.StatusNoContent, rec.Code)

	// Verify batch was canceled
	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, repository.BatchStatusCancelled, batch.Status)
}

func TestUploadDocumentBatch_Part1_Success(t *testing.T) {
	// Setup: Create testcontainers (database + localstack)
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)
	test.CreateAWSResources(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	// Create test client with unique ID
	clientID := fmt.Sprintf("test_client_part1_%d", time.Now().UnixNano())
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Part1",
	})
	require.NoError(t, err)

	// Create test ZIP file with 3 PDFs (test1.pdf, test2.pdf, test3.pdf)
	zipContent := createTestZIPFile(t, 3)

	// Create multipart form request
	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)

	part, err := writer.CreateFormFile("archive", "test-batch.zip")
	require.NoError(t, err)

	_, err = io.Copy(part, zipContent)
	require.NoError(t, err)

	err = writer.Close()
	require.NoError(t, err)

	// Test: Upload ZIP batch
	e := echo.New()
	rec := httptest.NewRecorder()
	ctx := e.NewContext(
		httptest.NewRequest(http.MethodPost, "/", body),
		rec,
	)
	ctx.Request().Header.Set("Content-Type", writer.FormDataContentType())

	err = cons.UploadDocumentBatch(ctx, clientID)
	require.NoError(t, err)

	// Verify: HTTP response
	assert.Equal(t, http.StatusAccepted, rec.Code)

	var response queryapi.BatchUploadResponse
	err = json.Unmarshal(rec.Body.Bytes(), &response)
	require.NoError(t, err)

	assert.NotEmpty(t, response.BatchId)
	assert.Equal(t, queryapi.BatchStatusProcessing, response.Status)
	assert.NotEmpty(t, response.StatusUrl)

	// Verify: Batch exists in database with storage metadata
	batch, err := cfg.GetDBQueries().GetBatchUploadWithStorage(t.Context(), &repository.GetBatchUploadWithStorageParams{
		ID:       uuid.UUID(response.BatchId),
		ClientID: clientID,
	})
	require.NoError(t, err)

	assert.Equal(t, "test-batch.zip", batch.OriginalFilename)
	assert.Equal(t, repository.BatchStatusProcessing, batch.Status)
	assert.NotNil(t, batch.ArchiveBucket)
	assert.NotEmpty(t, *batch.ArchiveBucket)
	assert.NotNil(t, batch.ArchiveKey)
	assert.NotEmpty(t, *batch.ArchiveKey)
	assert.Contains(t, *batch.ArchiveKey, clientID)
	assert.Contains(t, *batch.ArchiveKey, "test-batch.zip")
	assert.NotNil(t, batch.FileSizeBytes)
	assert.Greater(t, *batch.FileSizeBytes, int64(0))

	// Verify: ZIP file exists in S3
	s3Client := cfg.GetStoreClient()
	_, err = s3Client.HeadObject(t.Context(), &s3.HeadObjectInput{
		Bucket: batch.ArchiveBucket,
		Key:    batch.ArchiveKey,
	})
	require.NoError(t, err, "ZIP file should exist in S3")
}

func TestUploadDocumentBatch_Part1_InvalidFile(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	// Create test client with unique ID
	clientID := fmt.Sprintf("test_client_invalid_%d", time.Now().UnixNano())
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Invalid",
	})
	require.NoError(t, err)

	// Create invalid file (not ZIP)
	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)

	part, err := writer.CreateFormFile("archive", "test.pdf")
	require.NoError(t, err)

	_, err = part.Write([]byte("not a zip file"))
	require.NoError(t, err)

	err = writer.Close()
	require.NoError(t, err)

	// Test: Upload invalid file
	ctx := echo.New().NewContext(
		httptest.NewRequest(http.MethodPost, "/", body),
		httptest.NewRecorder(),
	)
	ctx.Request().Header.Set("Content-Type", writer.FormDataContentType())

	err = cons.UploadDocumentBatch(ctx, clientID)

	// Verify: Error response
	require.Error(t, err)
	httpErr, ok := err.(*echo.HTTPError)
	require.True(t, ok)
	assert.Equal(t, http.StatusBadRequest, httpErr.Code)
	assert.Contains(t, httpErr.Message, "invalid file type")
}

func TestUploadDocumentBatch_Part1_TooLarge(t *testing.T) {
	cfg := &ControllerConfig{}
	test.CreateDB(t, cfg)
	initializeTestConfig(t, cfg)

	svc := createControllerServices(cfg)
	cons := queryapi.NewControllers(svc, cfg)

	// Create test client with unique ID
	clientID := fmt.Sprintf("test_client_toolarge_%d", time.Now().UnixNano())
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client TooLarge",
	})
	require.NoError(t, err)

	// Create oversized ZIP file
	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)

	// Set a fake Content-Length header to simulate large file
	part, err := writer.CreateFormFile("archive", "huge.zip")
	require.NoError(t, err)

	// Write minimal content but set file size in header
	_, err = part.Write([]byte("PK")) // ZIP file header
	require.NoError(t, err)

	err = writer.Close()
	require.NoError(t, err)

	// Test: Upload with modified request to simulate large file
	req := httptest.NewRequest(http.MethodPost, "/", body)
	req.Header.Set("Content-Type", writer.FormDataContentType())

	// Override the file size in the form data
	ctx := echo.New().NewContext(req, httptest.NewRecorder())

	// Manually create form with large file size
	form, _ := ctx.MultipartForm()
	if len(form.File["archive"]) > 0 {
		// Simulate 101MB file
		form.File["archive"][0].Size = 101 * 1024 * 1024
	}

	err = cons.UploadDocumentBatch(ctx, clientID)

	// Verify: Error response
	require.Error(t, err)
	httpErr, ok := err.(*echo.HTTPError)
	require.True(t, ok)
	assert.Equal(t, http.StatusBadRequest, httpErr.Code)
	assert.Contains(t, httpErr.Message, "file too large")
}

// Helper function to create test ZIP with multiple PDFs
func createTestZIPFile(t *testing.T, pdfCount int) *bytes.Reader {
	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)

	// Create standard test PDF files
	standardFiles := []string{"test1.pdf", "test2.pdf", "test3.pdf"}

	// If requested count is <= 3, use standard names
	if pdfCount <= 3 {
		for i := 0; i < pdfCount; i++ {
			fileWriter, err := zipWriter.Create(standardFiles[i])
			require.NoError(t, err)

			// Write minimal PDF content
			pdfContent := fmt.Sprintf("%%PDF-1.4\n%%Test content for %s\n%%%%EOF", standardFiles[i])
			_, err = fileWriter.Write([]byte(pdfContent))
			require.NoError(t, err)
		}
	} else {
		// For more than 3 files, use numbered names
		for i := 0; i < pdfCount; i++ {
			filename := fmt.Sprintf("document_%d.pdf", i)
			fileWriter, err := zipWriter.Create(filename)
			require.NoError(t, err)

			// Write minimal PDF content
			pdfContent := fmt.Sprintf("%%PDF-1.4\n%%Document %d content\n%%%%EOF", i)
			_, err = fileWriter.Write([]byte(pdfContent))
			require.NoError(t, err)
		}
	}

	err := zipWriter.Close()
	require.NoError(t, err)

	return bytes.NewReader(buf.Bytes())
}