package api

import (
	"archive/zip"
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"log/slog"
	"path"
	"strings"

	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/document/batch"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/google/uuid"
)

// Config key constants for batch worker
const (
	ConfigKeyBatchService  = "batchService"
	ConfigKeyS3Client      = "s3Client"
	ConfigKeyBucket        = "bucket"
	ConfigKeyUploadHandler = "uploadHandler"
)

// acceptedExtensions defines the file extensions allowed in batch uploads.
// Macro-enabled extensions (.docm, .xlsm, .pptm) are intentionally excluded --
// documents with macros are rejected, so users must save as standard formats first.
var acceptedExtensions = map[string]bool{
	".pdf":  true,
	".tiff": true,
	".tif":  true,
	".jpeg": true,
	".jpg":  true,
	".png":  true,
	".bmp":  true,
	".docx": true,
	".xlsx": true,
	".pptx": true,
	".txt":  true,
	".eml":  true,
}

// isAcceptedExtension checks if a filename has a supported file extension.
func isAcceptedExtension(filename string) bool {
	ext := strings.ToLower(path.Ext(filename))
	return acceptedExtensions[ext]
}

// sanitizeZipPath cleans a path from a zip file to prevent path traversal attacks
// while preserving the internal folder structure.
// Returns empty string if the path is invalid or attempts traversal.
func sanitizeZipPath(p string) string {
	// Reject paths with backslashes (Windows path traversal)
	if strings.Contains(p, "\\") {
		return ""
	}

	// Clean the path to resolve . and .. elements
	cleaned := path.Clean(p)

	// Reject paths that try to escape (start with .. or /)
	if strings.HasPrefix(cleaned, "../") || strings.HasPrefix(cleaned, "/") || cleaned == ".." {
		return ""
	}

	// Check for net upward traversal
	// We need to ensure that we don't escape the root level
	parts := strings.Split(p, "/")
	depth := 0
	for _, part := range parts {
		if part == ".." {
			depth--
			if depth < 0 {
				// Trying to escape root
				return ""
			}
		} else if part != "." && part != "" {
			depth++
		}
	}

	// The cleaned path is safe and preserves internal structure
	return cleaned
}

// processBatchWork is the background worker function that processes batch uploads
func processBatchWork(ctx context.Context, logger *slog.Logger, config map[string]any) error {
	// Extract the batch service from config
	batchService, ok := config[ConfigKeyBatchService].(*batch.Service)
	if !ok {
		return fmt.Errorf("batch service not found in config")
	}

	// Extract the S3 client from config
	s3Client, ok := config[ConfigKeyS3Client].(*s3.Client)
	if !ok {
		return fmt.Errorf("S3 client not found in config")
	}

	// Extract the bucket name from config
	bucket, ok := config[ConfigKeyBucket].(string)
	if !ok {
		return fmt.Errorf("bucket not found in config")
	}

	// Extract the document upload handler from config
	uploadHandler, ok := config[ConfigKeyUploadHandler].(func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error)
	if !ok {
		return fmt.Errorf("upload handler not found in config")
	}

	// Query unprocessed batches
	batches, err := batchService.ListUnprocessed(ctx)
	if err != nil {
		logger.Error("Failed to query unprocessed batches", "error", err)
		return err
	}

	logger.Info("Background task executed", "unprocessed_batches", len(batches))

	// Process each batch
	for _, batchSummary := range batches {
		// Get full batch details with archive information
		batchDetails, err := batchService.Get(ctx, batchSummary.ClientID, batchSummary.ID)
		if err != nil {
			logger.Error("Failed to get batch details",
				"batch_id", batchSummary.ID,
				"client_id", batchSummary.ClientID,
				"error", err)
			continue
		}

		if err := processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails); err != nil {
			logger.Error("Failed to process batch",
				"batch_id", batchDetails.ID,
				"client_id", batchDetails.ClientID,
				"error", err)
			// Continue processing other batches even if one fails
			continue
		}
	}

	return nil
}

// processSingleBatch processes a single batch upload
func processSingleBatch(
	ctx context.Context,
	logger *slog.Logger,
	batchService *batch.Service,
	s3Client *s3.Client,
	bucket string,
	uploadHandler func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error,
	batchInfo *batch.BatchUploadDetails,
) error {
	logger.Info("Processing batch",
		"batch_id", batchInfo.ID,
		"client_id", batchInfo.ClientID,
		"filename", batchInfo.OriginalFilename)

	// Batch is already in 'processing' status by default when created

	// Download the zip file from S3
	zipData, err := downloadZipFromS3(ctx, s3Client, bucket, batchInfo.ArchiveKey)
	if err != nil {
		// Mark batch as failed if we can't download the zip
		_ = batchService.MarkFailed(ctx, batchInfo.ID)
		return fmt.Errorf("failed to download zip from S3: %w", err)
	}

	// Extract and process files from the zip
	processedCount := 0
	failedCount := 0
	invalidCount := 0

	// Create a temporary S3 directory for extracted files
	tempDir := fmt.Sprintf("%s_burst", batchInfo.ID.String())
	basePath := path.Dir(batchInfo.ArchiveKey)
	extractPath := path.Join(basePath, tempDir)

	// Open the zip archive
	zipReader, err := zip.NewReader(bytes.NewReader(zipData), int64(len(zipData)))
	if err != nil {
		_ = batchService.MarkFailed(ctx, batchInfo.ID)
		return fmt.Errorf("failed to open zip archive: %w", err)
	}

	// Process each file in the zip
	for _, file := range zipReader.File {
		processed, failed, invalid := processZipFile(ctx, logger, batchService, s3Client, bucket, extractPath, uploadHandler, batchInfo, file)
		processedCount += processed
		failedCount += failed
		invalidCount += invalid
	}

	// Update batch progress - convert safely to int32
	const maxInt32 = 2147483647
	var processedInt32, failedInt32, invalidInt32 int32
	if processedCount <= maxInt32 {
		processedInt32 = int32(processedCount) // #nosec G115 - checked bounds above
	} else {
		processedInt32 = maxInt32 // Max int32 value
	}
	if failedCount <= maxInt32 {
		failedInt32 = int32(failedCount) // #nosec G115 - checked bounds above
	} else {
		failedInt32 = maxInt32 // Max int32 value
	}
	if invalidCount <= maxInt32 {
		invalidInt32 = int32(invalidCount) // #nosec G115 - checked bounds above
	} else {
		invalidInt32 = maxInt32 // Max int32 value
	}

	// Set total document count now that all files have been extracted
	totalInt32 := processedInt32 + failedInt32 + invalidInt32
	if err := batchService.SetTotalDocuments(ctx, batchInfo.ID, totalInt32); err != nil {
		logger.Error("Failed to set batch total documents", "batch_id", batchInfo.ID, "error", err)
	}

	if err := batchService.UpdateProgress(ctx, batchInfo.ClientID, batchInfo.ID,
		processedInt32, failedInt32, invalidInt32); err != nil {
		logger.Error("Failed to update batch progress", "batch_id", batchInfo.ID, "error", err)
	}

	// Mark batch as completed if any files were processed successfully
	totalProcessed := processedCount + failedCount + invalidCount
	if totalProcessed > 0 {
		if processedCount > 0 {
			// Mark as completed if at least one document was processed successfully
			if err := batchService.MarkCompleted(ctx, batchInfo.ID); err != nil {
				logger.Error("Failed to mark batch as completed", "batch_id", batchInfo.ID, "error", err)
			}
		} else {
			// Mark as failed only if no documents were processed successfully
			if err := batchService.MarkFailed(ctx, batchInfo.ID); err != nil {
				logger.Error("Failed to mark batch as failed", "batch_id", batchInfo.ID, "error", err)
			}
		}
	}

	logger.Info("Batch processing completed",
		"batch_id", batchInfo.ID,
		"processed", processedCount,
		"failed", failedCount,
		"invalid", invalidCount)

	return nil
}

// downloadZipFromS3 downloads a zip file from S3
func downloadZipFromS3(ctx context.Context, s3Client *s3.Client, bucket, key string) ([]byte, error) {
	result, err := s3Client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return nil, fmt.Errorf("failed to get object from S3: %w", err)
	}
	defer result.Body.Close()

	data, err := io.ReadAll(result.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read S3 object body: %w", err)
	}

	return data, nil
}

// uploadToS3 uploads data to S3
func uploadToS3(ctx context.Context, s3Client *s3.Client, bucket, key string, data []byte) error {
	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
		Body:   bytes.NewReader(data),
	})
	if err != nil {
		return fmt.Errorf("failed to put object to S3: %w", err)
	}
	return nil
}

// processZipFile processes a single file from the zip archive
// Returns (processedCount, failedCount, invalidCount)
func processZipFile(
	ctx context.Context,
	logger *slog.Logger,
	batchService *batch.Service,
	s3Client *s3.Client,
	bucket, extractPath string,
	uploadHandler func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error,
	batchInfo *batch.BatchUploadDetails,
	file *zip.File,
) (int, int, int) {
	if file.FileInfo().IsDir() {
		return 0, 0, 0
	}

	// Skip system files and hidden files (but allow relative path prefixes like ./file.pdf)
	filename := file.Name
	if strings.HasPrefix(filename, "__MACOSX/") {
		return 0, 0, 0
	}

	// Skip hidden files (files that start with . but are not relative paths like ./file)
	if strings.HasPrefix(filename, ".") && !strings.HasPrefix(filename, "./") {
		return 0, 0, 0
	}

	// Security: Clean filename to prevent path traversal attacks while preserving internal structure
	filename = sanitizeZipPath(filename)
	if filename == "" {
		logger.Warn("Skipping file with invalid path", "original", file.Name)
		return 0, 0, 0
	}

	// Extract file content
	fileReader, err := file.Open()
	if err != nil {
		logger.Error("Failed to open file in zip", "filename", filename, "error", err)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		errMsg := err.Error()
		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedOpen, &errMsg, nil)
		return 0, 1, 0
	}

	fileData, err := io.ReadAll(fileReader)
	fileReader.Close()
	if err != nil {
		logger.Error("Failed to read file from zip", "filename", filename, "error", err)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		errMsg := err.Error()
		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedRead, &errMsg, nil)
		return 0, 1, 0
	}

	// Store extracted file in S3 temporary directory (filename now includes sanitized path)
	extractedKey := extractPath + "/" + filename
	if err := uploadToS3(ctx, s3Client, bucket, extractedKey, fileData); err != nil {
		logger.Error("Failed to upload extracted file to S3", "filename", filename, "error", err)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		errMsg := err.Error()
		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedS3Upload, &errMsg, nil)
		return 0, 1, 0
	}

	// Check if file type is supported
	if !isAcceptedExtension(filename) {
		logger.Warn("Invalid file type", "filename", filename)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusInvalidType, nil, nil)
		return 0, 0, 1
	}

	// Call the document upload handler with batch ID
	if err := uploadHandler(ctx, batchInfo.ClientID, bytes.NewReader(fileData), filename, &batchInfo.ID); err != nil {
		logger.Error("Failed to upload document", "filename", filename, "error", err)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		errMsg := err.Error()
		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedUpload, &errMsg, nil)
		return 0, 1, 0
	}

	// Check for duplicate by computing file hash and looking up existing document
	hash := sha256.Sum256(fileData)
	hashStr := hex.EncodeToString(hash[:])
	existingID, _ := batchService.CheckDuplicate(ctx, batchInfo.ClientID, hashStr)
	if existingID != nil {
		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusDuplicate, nil, existingID)
	} else {
		_ = batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusSubmitted, nil, nil)
	}

	logger.Info("Successfully processed file", "filename", filename, "batch_id", batchInfo.ID)
	return 1, 0, 0
}