package loader

import (
	"context"
	"errors"
	"sort"
	"strings"
	"testing"

	"github.com/pingcap/tidb/pkg/executor/importer"
	"github.com/pingcap/tidb/pkg/lightning/mydump"
	"github.com/pingcap/tidb/pkg/objstore/storeapi"
	"github.com/stretchr/testify/require"
)

// mockFile represents a file in the mock storage.
type mockFile struct {
	size int64
}

// mockStorage is a simple in-memory ExternalStorage that handles ObjPrefix
// the same way as S3: prefix-matching against the full path, not just the base name.
// This differs from MemStorage which only checks path.Base(fileName).
type mockStorage struct {
	storeapi.Storage // embed to satisfy interface; unused methods will panic
	files                   map[string]mockFile
}

func newMockStorage(files map[string]mockFile) *mockStorage {
	return &mockStorage{files: files}
}

func (m *mockStorage) WalkDir(_ context.Context, opt *storeapi.WalkOption, fn func(string, int64) error) error {
	names := make([]string, 0, len(m.files))
	for name := range m.files {
		if opt != nil && len(opt.ObjPrefix) > 0 {
			if !strings.HasPrefix(name, opt.ObjPrefix) {
				continue
			}
		}
		names = append(names, name)
	}
	sort.Strings(names)
	for _, name := range names {
		if err := fn(name, m.files[name].size); err != nil {
			return err
		}
	}
	return nil
}

// withMockStore replaces getStoreFn with one that returns the given mock storage,
// and registers a cleanup to restore the original via t.Cleanup.
func withMockStore(t *testing.T, ms *mockStorage) {
	t.Helper()
	orig := getStoreFn
	getStoreFn = func(_ context.Context, _ string) (storeapi.Storage, error) {
		return ms, nil
	}
	t.Cleanup(func() { getStoreFn = orig })
}

// withFailingStore replaces getStoreFn with one that always returns an error.
func withFailingStore(t *testing.T, err error) {
	t.Helper()
	orig := getStoreFn
	getStoreFn = func(_ context.Context, _ string) (storeapi.Storage, error) {
		return nil, err
	}
	t.Cleanup(func() { getStoreFn = orig })
}

// ---------------------------------------------------------------------------
// Tests for ScanDataFiles – single file path (no glob)
// ---------------------------------------------------------------------------

func TestScanDataFiles_SingleFile_Found_CSV(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{
		"path/data.csv": {size: 1024},
	})
	withMockStore(t, ms)

	files, totalSize, err := ScanDataFiles(context.Background(), "s3://bucket/path/data.csv", 0)
	require.NoError(t, err)
	require.Len(t, files, 1)
	require.Equal(t, int64(1024), totalSize)
	require.Equal(t, "path/data.csv", files[0].Path)
	require.Equal(t, int64(1024), files[0].FileSize)
	require.Equal(t, mydump.SourceTypeCSV, files[0].Type)
	require.Equal(t, mydump.CompressionNone, files[0].Compression)
}

func TestScanDataFiles_SingleFile_Found_SQL(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{
		"schema.sql": {size: 256},
	})
	withMockStore(t, ms)

	files, totalSize, err := ScanDataFiles(context.Background(), "s3://bucket/schema.sql", 0)
	require.NoError(t, err)
	require.Len(t, files, 1)
	require.Equal(t, int64(256), totalSize)
	require.Equal(t, mydump.SourceTypeSQL, files[0].Type)
}

func TestScanDataFiles_SingleFile_Found_Parquet(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{
		"data.parquet": {size: 4096},
	})
	withMockStore(t, ms)

	files, totalSize, err := ScanDataFiles(context.Background(), "s3://bucket/data.parquet", 0)
	require.NoError(t, err)
	require.Len(t, files, 1)
	require.Equal(t, int64(4096), totalSize)
	require.Equal(t, mydump.SourceTypeParquet, files[0].Type)
}

func TestScanDataFiles_SingleFile_Found_GzipCompressed(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{
		"data.csv.gz": {size: 512},
	})
	withMockStore(t, ms)

	files, _, err := ScanDataFiles(context.Background(), "s3://bucket/data.csv.gz", 0)
	require.NoError(t, err)
	require.Len(t, files, 1)
	require.Equal(t, mydump.SourceTypeCSV, files[0].Type)
	require.Equal(t, mydump.CompressionGZ, files[0].Compression)
}

func TestScanDataFiles_SingleFile_NotFound(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{
		"other.csv": {size: 100},
	})
	withMockStore(t, ms)

	files, totalSize, err := ScanDataFiles(context.Background(), "s3://bucket/missing.csv", 0)
	require.Error(t, err)
	require.True(t, errors.Is(err, ErrFileNotFound))
	require.Contains(t, err.Error(), "missing.csv")
	require.Nil(t, files)
	require.Equal(t, int64(0), totalSize)
}

func TestScanDataFiles_SingleFile_EmptyStorage(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{})
	withMockStore(t, ms)

	_, _, err := ScanDataFiles(context.Background(), "s3://bucket/data.csv", 0)
	require.Error(t, err)
	require.True(t, errors.Is(err, ErrFileNotFound))
}

func TestScanDataFiles_SingleFile_PrefixMatchDoesNotCount(t *testing.T) {
	// A file "data.csv.bak" should NOT match when scanning for "data.csv".
	// WalkDir returns it because of prefix, but the exact-match filter should exclude it.
	ms := newMockStorage(map[string]mockFile{
		"data.csv.bak": {size: 999},
	})
	withMockStore(t, ms)

	_, _, err := ScanDataFiles(context.Background(), "s3://bucket/data.csv", 0)
	require.Error(t, err)
	require.True(t, errors.Is(err, ErrFileNotFound))
}

// ---------------------------------------------------------------------------
// Tests for ScanDataFiles – glob patterns
// ---------------------------------------------------------------------------

func TestScanDataFiles_Glob_MatchesMultipleFiles(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{
		"data001.csv": {size: 100},
		"data002.csv": {size: 200},
		"data003.csv": {size: 300},
		"other.sql":   {size: 50},
	})
	withMockStore(t, ms)

	files, totalSize, err := ScanDataFiles(context.Background(), "s3://bucket/data*.csv", 0)
	require.NoError(t, err)
	require.Len(t, files, 3)
	require.Equal(t, int64(600), totalSize)
	for _, f := range files {
		require.Equal(t, mydump.SourceTypeCSV, f.Type)
	}
}

func TestScanDataFiles_Glob_NoMatches(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{
		"other.csv": {size: 100},
	})
	withMockStore(t, ms)

	files, totalSize, err := ScanDataFiles(context.Background(), "s3://bucket/data*.csv", 0)
	require.NoError(t, err)
	require.Len(t, files, 0)
	require.Equal(t, int64(0), totalSize)
}

func TestScanDataFiles_Glob_MaxScanFilesLimit(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{
		"data001.csv": {size: 100},
		"data002.csv": {size: 200},
		"data003.csv": {size: 300},
		"data004.csv": {size: 400},
		"data005.csv": {size: 500},
	})
	withMockStore(t, ms)

	// maxScanFiles=3 means WalkDir stops after collecting 3 files
	files, _, err := ScanDataFiles(context.Background(), "s3://bucket/data*.csv", 3)
	require.NoError(t, err)
	// Should have at most 3 files (the first 3 in sorted order)
	require.LessOrEqual(t, len(files), 3)
}

func TestScanDataFiles_Glob_BracketPattern(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{
		"data_a.csv": {size: 100},
		"data_b.csv": {size: 200},
		"data_c.csv": {size: 300},
	})
	withMockStore(t, ms)

	files, totalSize, err := ScanDataFiles(context.Background(), "s3://bucket/data_[ab].csv", 0)
	require.NoError(t, err)
	require.Len(t, files, 2)
	require.Equal(t, int64(300), totalSize)
}

// ---------------------------------------------------------------------------
// Tests for ScanDataFiles – compression detection
// ---------------------------------------------------------------------------

func TestScanDataFiles_Compression_Zstd(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{
		"data.csv.zst": {size: 1024},
	})
	withMockStore(t, ms)

	files, _, err := ScanDataFiles(context.Background(), "s3://bucket/data.csv.zst", 0)
	require.NoError(t, err)
	require.Len(t, files, 1)
	require.Equal(t, mydump.SourceTypeCSV, files[0].Type)
	require.Equal(t, mydump.CompressionZStd, files[0].Compression)
}

func TestScanDataFiles_Compression_Snappy(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{
		"data.csv.snappy": {size: 1024},
	})
	withMockStore(t, ms)

	files, _, err := ScanDataFiles(context.Background(), "s3://bucket/data.csv.snappy", 0)
	require.NoError(t, err)
	require.Len(t, files, 1)
	require.Equal(t, mydump.SourceTypeCSV, files[0].Type)
	require.Equal(t, mydump.CompressionSnappy, files[0].Compression)
}

// ---------------------------------------------------------------------------
// Tests for ScanDataFiles – error cases
// ---------------------------------------------------------------------------

func TestScanDataFiles_InvalidURL(t *testing.T) {
	_, _, err := ScanDataFiles(context.Background(), "://invalid", 0)
	require.Error(t, err)
}

func TestScanDataFiles_LocalFileRejected(t *testing.T) {
	_, _, err := ScanDataFiles(context.Background(), "local:///tmp/data.csv", 0)
	require.Error(t, err)
	require.Contains(t, err.Error(), "local file is not supported")
}

func TestScanDataFiles_StoreCreationError(t *testing.T) {
	storeErr := errors.New("access denied")
	withFailingStore(t, storeErr)

	_, _, err := ScanDataFiles(context.Background(), "s3://bucket/data.csv", 0)
	require.Error(t, err)
	require.Equal(t, storeErr, err)
}

func TestScanDataFiles_ContextCancelled(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel() // cancel immediately

	ms := newMockStorage(map[string]mockFile{
		"data.csv": {size: 100},
	})
	withMockStore(t, ms)

	// This should complete (not hang). Result may be success or error.
	_, _, _ = ScanDataFiles(ctx, "s3://bucket/data.csv", 0)
}

// ---------------------------------------------------------------------------
// Tests for parseFileType (internal helper)
// ---------------------------------------------------------------------------

func TestParseFileType(t *testing.T) {
	tests := []struct {
		path     string
		expected string
	}{
		{"data.csv", importer.DataFormatCSV},
		{"data.CSV", importer.DataFormatCSV},
		{"data.sql", importer.DataFormatSQL},
		{"data.SQL", importer.DataFormatSQL},
		{"data.parquet", importer.DataFormatParquet},
		{"data.PARQUET", importer.DataFormatParquet},
		{"data.tsv", importer.DataFormatCSV},            // default to csv
		{"data", importer.DataFormatCSV},                // no extension defaults to csv
		{"data.csv.gz", importer.DataFormatCSV},         // strip gz, then detect csv
		{"data.csv.gzip", importer.DataFormatCSV},       // strip gzip
		{"data.csv.zstd", importer.DataFormatCSV},       // strip zstd
		{"data.csv.zst", importer.DataFormatCSV},        // strip zst
		{"data.csv.snappy", importer.DataFormatCSV},     // strip snappy
		{"data.sql.gz", importer.DataFormatSQL},         // strip gz, then detect sql
		{"data.parquet.gz", importer.DataFormatParquet}, // strip gz, then detect parquet
	}
	for _, tt := range tests {
		t.Run(tt.path, func(t *testing.T) {
			result := parseFileType(tt.path)
			require.Equal(t, tt.expected, result)
		})
	}
}

// ---------------------------------------------------------------------------
// Tests for getSourceType (internal helper)
// ---------------------------------------------------------------------------

func TestGetSourceType(t *testing.T) {
	require.Equal(t, mydump.SourceTypeParquet, getSourceType(importer.DataFormatParquet))
	require.Equal(t, mydump.SourceTypeCSV, getSourceType(importer.DataFormatCSV))
	require.Equal(t, mydump.SourceTypeCSV, getSourceType(importer.DataFormatDelimitedData))
	require.Equal(t, mydump.SourceTypeSQL, getSourceType(importer.DataFormatSQL))
	require.Equal(t, mydump.SourceTypeSQL, getSourceType("unknown"))
}

// ---------------------------------------------------------------------------
// Tests for ErrFileNotFound sentinel error
// ---------------------------------------------------------------------------

func TestErrFileNotFound_Wrapping(t *testing.T) {
	ms := newMockStorage(map[string]mockFile{})
	withMockStore(t, ms)

	_, _, err := ScanDataFiles(context.Background(), "s3://bucket/missing.csv", 0)
	require.Error(t, err)
	require.True(t, errors.Is(err, ErrFileNotFound))
	require.Contains(t, err.Error(), "file not found")
	require.Contains(t, err.Error(), "missing.csv")
}