package import_data

import (
	"database/sql"
	"fmt"
	"strings"
	"time"

	"github.com/go-openapi/swag"
	"github.com/tidbcloud/dataflow-service-e2e/internal/config"
	. "github.com/tidbcloud/dataflow-service-e2e/internal/model"
	"github.com/tidbcloud/dataflow-service-e2e/internal/sample_data"
	import_service "github.com/tidbcloud/dataflow-service-e2e/internal/tidbcloud/dedicated/consoleapi/import/client/import_service"
)

// Common test constants for import E2E tests
const (
	// Timeout constants
	DefaultImportTimeout = 30 * time.Minute

	// Retry constants
	DefaultRetryAttempts = 3
	DefaultRetryInterval = time.Minute

	// Data format: SQL (dumpling default format)
	DataFormatSQL int64 = 1

	// Data format: aurora snapshot (parquet dumped from aurora)
	DataFormatAuroraSnapshot int64 = 2

	// Data format: custom csv
	DataFormatCSV int64 = 3

	// Data format: parquet (not aurora)
	DataFormatParquet int64 = 4

	// Data source types
	DataSourceTypeS3     int64 = 1
	DataSourceTypeGCS    int64 = 2
	DataSourceTypeLocal  int64 = 3
	DataSourceTypeAzBlob int64 = 4

	// Imported database/table for data validation
	ImportedDatabase = "bikeshare"
	ImportedTable    = "trips"
)

// GenerateImportName creates a unique import name for test purposes.
func GenerateImportName(prefix string) string {
	return fmt.Sprintf("%s-%d", prefix, time.Now().Unix())
}

// BuildImportBody builds an import request body for the given cloud provider.
func BuildImportBody(name string, cluster *TiDBDedicatedEntity) import_service.CreateImportBody {
	switch cluster.CloudProvider {
	case GCP:
		return buildGCSImportBody(name, cluster)
	case AZURE:
		return buildAzBlobImportBody(name, cluster)
	default:
		return buildS3ImportBody(name, cluster)
	}
}

func buildS3ImportBody(name string, cluster *TiDBDedicatedEntity) import_service.CreateImportBody {
	sampleData := config.GlobalConfig.ImportSampleData
	return import_service.CreateImportBody{
		Name:           name,
		SourceURL:      sampleData.S3.SourceURL,
		DataFormat:     swag.Int64(DataFormatSQL),
		DataSourceType: swag.Int64(DataSourceTypeS3),
		Region:         swag.String(cluster.Region),
		AwsRoleArn:     sampleData.S3.RoleARN,
	}
}

func buildGCSImportBody(name string, cluster *TiDBDedicatedEntity) import_service.CreateImportBody {
	sampleData := config.GlobalConfig.ImportSampleData
	return import_service.CreateImportBody{
		Name:           name,
		SourceURL:      sampleData.GCS.SourceURL,
		DataFormat:     swag.Int64(DataFormatSQL),
		DataSourceType: swag.Int64(DataSourceTypeGCS),
		Region:         swag.String(cluster.Region),
	}
}

func buildAzBlobImportBody(name string, cluster *TiDBDedicatedEntity) import_service.CreateImportBody {
	sampleData := config.GlobalConfig.ImportSampleData
	return import_service.CreateImportBody{
		Name:           name,
		SourceURL:      sampleData.AzureBlob.SourceURL,
		DataFormat:     swag.Int64(DataFormatSQL),
		DataSourceType: swag.Int64(DataSourceTypeAzBlob),
		Region:         swag.String(cluster.Region),
		AzblobSasToken: sampleData.AzureBlob.SASToken,
	}
}

// legacyDataFormat maps the canonical importData dataFormat to the legacy
// dedicated console import int64 data_format enum.
func legacyDataFormat(dataFormat string) (int64, error) {
	switch dataFormat {
	case config.DataFormatSQL:
		return DataFormatSQL, nil
	case config.DataFormatAuroraSnapshot:
		return DataFormatAuroraSnapshot, nil
	case config.DataFormatCSV:
		return DataFormatCSV, nil
	case config.DataFormatParquet:
		return DataFormatParquet, nil
	default:
		return 0, fmt.Errorf("unsupported dedicated import data format %q", dataFormat)
	}
}

// BuildImportBodyWithData builds a dedicated import request from a configured
// importData item (S3 source with an explicit data format).
func BuildImportBodyWithData(name string, cluster *TiDBDedicatedEntity, item config.ImportDataItem) (import_service.CreateImportBody, error) {
	dataFormat, err := legacyDataFormat(item.DataFormat)
	if err != nil {
		return import_service.CreateImportBody{}, err
	}
	if strings.TrimSpace(item.SourceURL) == "" || strings.TrimSpace(item.RoleARN) == "" {
		return import_service.CreateImportBody{}, fmt.Errorf("import data item %s is missing sourceURL or roleARN", item.DataFormat)
	}
	return import_service.CreateImportBody{
		Name:           name,
		SourceURL:      item.SourceURL,
		DataFormat:     swag.Int64(dataFormat),
		DataSourceType: swag.Int64(DataSourceTypeS3),
		Region:         swag.String(cluster.Region),
		AwsRoleArn:     item.RoleARN,
	}, nil
}

// verifyImportedDataIn connects to the TiDB cluster and checks that the given
// database contains the expected rows for every declared table (expectedRows 0
// = at least one row, positive = exact count).
func verifyImportedDataIn(cluster *TiDBDedicatedEntity, database string, tables []config.ImportDataTableName) error {
	db, err := sample_data.Connect2TiDB(
		cluster.User, cluster.Password,
		cluster.Host, cluster.Port,
		cluster.CACert, database,
	)
	if err != nil {
		return fmt.Errorf("connect to TiDB failed: %w", err)
	}
	defer func(db *sql.DB) {
		_ = db.Close()
	}(db)

	for _, table := range tables {
		var count int64
		err = db.QueryRow(fmt.Sprintf("SELECT COUNT(*) FROM `%s`", table.Table)).Scan(&count)
		if err != nil {
			return fmt.Errorf("query %s.%s failed: %w", database, table.Table, err)
		}
		if table.ExpectedRows > 0 && count != table.ExpectedRows {
			return fmt.Errorf("%s.%s contains %d rows, want %d", database, table.Table, count, table.ExpectedRows)
		}
		if count == 0 {
			return fmt.Errorf("expected rows in %s.%s but got 0", database, table.Table)
		}
	}
	return nil
}