package networktask

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"strings"
	"testing"

	"github.com/aws/aws-sdk-go-v2/aws"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/stretchr/testify/assert"
	"github.com/tidbcloud/metering/pkg/common/config"
	"github.com/tidbcloud/metering/pkg/common/constants"
	"github.com/tidbcloud/metering/pkg/common/models"
	"github.com/tidbcloud/metering/pkg/networkmetering"
	"github.com/tidbcloud/metering/pkg/repoter"
	"go.uber.org/zap"
)

func TestNetWorkMeteringFromAWS(t *testing.T) {
	//12-21 00:00:00 UTC
	ts := int64(1734739200)
	conf := config.ServiceConfig{
		Vendor: constants.VendorAWS,
		Env:    constants.EnvDev,
		Secret: &config.Secret{
			ApiGatewayKey: "test",
		},
		TrafficMeteringConfig: &config.TrafficMeteringConfig{
			Provider: config.TrafficMeteringProvider{
				AWS: &config.TrafficMeteringProviderAWSConfig{
					BillingProvider: []config.BillingProviderAWSConfig{
						{
							BucketName: "costmetetringstoragedev",
							Account:    "385595570414",
							Perfix:     "dataplane-cost-storage/test-cost/",
							Region:     "us-west-2",
							SourceType: constants.AWSBillingSourceTypeParquet,
						},
					},
					StoreProvider: config.StoreProviderAWSConfig{
						BucketName: "costmetetringstoragedev",
						Region:     "us-west-2",
						Account:    "385595570414",
						Perfix:     "dataplane-cost-storage/test-cost/billingresult",
					},
				},
			},
			CachePath: "cachetest",
		},
	}
	// set aws profile if needed
	os.Setenv("AWS_PROFILE", "dbaas-dev")
	cli := NewAWSS3(conf.TrafficMeteringConfig.Provider.AWS.StoreProvider.BucketName, "us-west-2")
	// clean test data
	cleanTestDataAWS(cli)
	defer cleanTestDataAWS(cli)
	//prepare test data
	prepareTestDataAWS(cli)
	logger, _ := zap.NewDevelopment()
	metering, err := networkmetering.NewNetworkMetering(&conf, logger)
	assert.Nil(t, err)
	//set mock gateway because we can't control the api gateway in this test
	metering.SetGateway(&mockAPIGateway{})
	path, ok, err := metering.Metering(context.Background(), ts)
	assert.Nil(t, err)
	assert.True(t, ok)
	assert.Equal(t, "dataplane-cost-storage/test-cost/billingresult/20241221/20241221.db", path)
	//compare the result with the expected result
	ok, err = compareResult(t, "./cachetest/resultcache/20241221.db", "../../../testdata/awsbillingfile/20241221_normal.db")
	assert.Nil(t, err)
	assert.True(t, ok)
}

func TestNetWorkMeteringFromAWSRerun(t *testing.T) {
	//12-21 00:00:00 UTC
	ts := int64(1734739200)
	conf := config.ServiceConfig{
		Vendor: constants.VendorAWS,
		Env:    constants.EnvDev,
		Secret: &config.Secret{
			ApiGatewayKey: "test",
		},
		TrafficMeteringConfig: &config.TrafficMeteringConfig{
			Provider: config.TrafficMeteringProvider{
				AWS: &config.TrafficMeteringProviderAWSConfig{
					BillingProvider: []config.BillingProviderAWSConfig{
						{
							BucketName: "costmetetringstoragedev",
							Account:    "385595570414",
							Perfix:     "dataplane-cost-storage/test-cost/",
							Region:     "us-west-2",
							SourceType: constants.AWSBillingSourceTypeParquet,
						},
					},
					StoreProvider: config.StoreProviderAWSConfig{
						BucketName: "costmetetringstoragedev",
						Region:     "us-west-2",
						Account:    "385595570414",
						Perfix:     "dataplane-cost-storage/test-cost/rerunbillingresult",
					},
				},
			},
			CachePath: "cachetest",
		},
	}
	// set aws profile if needed
	os.Setenv("AWS_PROFILE", "dbaas-dev")
	cli := NewAWSS3(conf.TrafficMeteringConfig.Provider.AWS.StoreProvider.BucketName, "us-west-2")
	// clean test data
	cleanTestDataAWS(cli)
	defer cleanTestDataAWS(cli)
	//prepare test data
	prepareTestDataAWS(cli)
	logger, _ := zap.NewDevelopment()
	metering, err := networkmetering.NewNetworkMetering(&conf, logger)
	assert.Nil(t, err)
	//set mock gateway because we can't control the api gateway in this test
	metering.SetGateway(&mockAPIGateway{})
	path, ok, err := metering.Metering(context.Background(), ts)
	assert.Nil(t, err)
	assert.True(t, ok)
	assert.Equal(t, "dataplane-cost-storage/test-cost/rerunbillingresult/20241221/20241221.db", path)
	//compare the result with the expected result
	ok, err = compareResult(t, "./cachetest/resultcache/20241221.db", "../../../testdata/awsbillingfile/20241221_normal.db")
	assert.Nil(t, err)
	assert.True(t, ok)
}

func TestNetWorkMeteringFromAWSRepoter(t *testing.T) {
	conf := config.ServiceConfig{
		Vendor: constants.VendorAWS,
		Env:    constants.EnvDev,
		Region: "us-west-2",
		Secret: &config.Secret{
			ApiGatewayKey: "test",
		},
		TrafficMeteringConfig: &config.TrafficMeteringConfig{
			Provider: config.TrafficMeteringProvider{
				AWS: &config.TrafficMeteringProviderAWSConfig{
					BillingProvider: []config.BillingProviderAWSConfig{
						{
							BucketName: "costmetetringstoragedev",
							Account:    "385595570414",
							Perfix:     "dataplane-cost-storage/test-cost/",
							Region:     "us-west-2",
							SourceType: constants.AWSBillingSourceTypeParquet,
						},
					},
					StoreProvider: config.StoreProviderAWSConfig{
						BucketName: "costmetetringstoragedev",
						Region:     "us-west-2",
						Account:    "385595570414",
						Perfix:     "dataplane-cost-storage/test-cost/repoterbillingresult",
					},
				},
			},
			CachePath: "cachetest",
		},
	}
	// set aws profile if needed
	os.Setenv("AWS_PROFILE", "dbaas-dev")
	cli := NewAWSS3(conf.TrafficMeteringConfig.Provider.AWS.StoreProvider.BucketName, "us-west-2")
	// clean test data
	deleteTestDataAWSRepoter(cli)
	defer deleteTestDataAWSRepoter(cli)
	//prepare test data
	prepareTestDataAWSRepoter(cli)
	logger, _ := zap.NewDevelopment()
	repoter, err := repoter.NewBillingRepoter(&conf, logger)
	assert.Nil(t, err)
	//set mock sender because we can't control the billing server in this test
	s := &MockMeasureSender{t: t}
	repoter.SetSender(s)
	err = repoter.Report(context.Background(), "dataplane-cost-storage/test-cost/repoterbillingresult/20241221/20241221.db")
	assert.Nil(t, err)
	//compare the result with the expected result
	ok, err := compareMeasure(t, s.GetSended(), "../../../testdata/awsbillingfile/20241221_normal.db")
	assert.Nil(t, err)
	assert.True(t, ok)
	//get checkpoint
	checkpoint, err := cli.DownloadObjects(context.Background(), "dataplane-cost-storage/test-cost/repoterbillingresult/20241221/checkpoint.json")
	assert.Nil(t, err)
	assert.Equal(t, string(checkpoint), "{\"timestamp\":1734869187}\n")
}

func compareMeasure(t *testing.T, sended []models.BillingReport, expectedfile string) (bool, error) {
	expectedDb, err := loadResult(expectedfile)
	if err != nil {
		panic(err)
	}
	var reportresExpected []*models.BillingReport
	err = expectedDb.Find(&reportresExpected).Error
	if err != nil {
		return false, err
	}
	//make sure the sended measure is in the expected measure
	for _, measure := range sended {
		found := false
		for _, expected := range reportresExpected {
			if measure.TenantID == expected.TenantID && measure.CostID == expected.CostID && measure.Cost == expected.Cost && measure.CostTimestamp == expected.CostTimestamp && measure.Unit == expected.Unit {
				found = true
				break
			}
		}
		assert.True(t, found, "measure %+v not found in expected measure", measure)
	}
	//make sure the expected measure is in the sended measure
	for _, measure := range reportresExpected {
		found := false
		for _, sended := range sended {
			if measure.TenantID == sended.TenantID && measure.CostID == sended.CostID && measure.Cost == sended.Cost && measure.CostTimestamp == sended.CostTimestamp && measure.Unit == sended.Unit {
				found = true
				break
			}
		}
		assert.True(t, found, "expected measure %+v not found in sended measure", measure)
	}
	return true, nil
}

func cleanTestDataAWS(cli *AWSS3) {
	cli.Delete(context.Background(), "dataplane-cost-storage/test-cost/20241201-20250101/dev-billing-Manifest.json")
	cli.Delete(context.Background(), "dataplane-cost-storage/test-cost/20241201-20250101/dev-billing-00001.snappy.parquet")
	cli.Delete(context.Background(), "dataplane-cost-storage/test-cost/rerunbillingresult/20241221/20241221.db")
	cli.Delete(context.Background(), "dataplane-cost-storage/test-cost/billingresult/20241221/20241221.db")
}

func prepareTestDataAWS(cli *AWSS3) {
	cli.Upload(context.Background(), "dataplane-cost-storage/test-cost/20241201-20250101/dev-billing-Manifest.json", "../../../testdata/awsbillingfile/ft_billing-Manifest.json")
	cli.Upload(context.Background(), "dataplane-cost-storage/test-cost/20241201-20250101/dev-billing-00001.snappy.parquet", "../../../testdata/awsbillingfile/ft_billing.snappy.parquet")
	cli.Upload(context.Background(), "dataplane-cost-storage/test-cost/rerunbillingresult/20241221/20241221.db", "../../../testdata/awsbillingfile/20241221_rerun.db")
}

func prepareTestDataAWSRepoter(cli *AWSS3) {
	cli.Upload(context.Background(), "dataplane-cost-storage/test-cost/repoterbillingresult/20241221/20241221.db", "../../../testdata/awsbillingfile/20241221_normal.db")
}

func deleteTestDataAWSRepoter(cli *AWSS3) {
	cli.Delete(context.Background(), "dataplane-cost-storage/test-cost/repoterbillingresult/20241221/20241221.db")
	cli.Delete(context.Background(), "dataplane-cost-storage/test-cost/repoterbillingresult/20241221/checkpoint.json")
}

type AWSS3 struct {
	bucketName string
	region     string
	cli        *s3.Client
}

func NewAWSS3(bucketName string, region string) *AWSS3 {
	cfg, err := awsconfig.LoadDefaultConfig(context.TODO())
	if err != nil {
		panic(fmt.Errorf("unable to get aws config: %v", err))
	}
	cfg.Region = region
	return &AWSS3{
		bucketName: bucketName,
		region:     region,
		cli:        s3.NewFromConfig(cfg),
	}
}

func (a *AWSS3) Upload(ctx context.Context, key string, path string) error {
	file, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("failed to open file %s: %w", path, err)
	}
	defer file.Close()
	p, err := io.ReadAll(file)
	if err != nil {
		return fmt.Errorf("failed to read file %s: %w", path, err)
	}
	_, err = a.cli.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(a.bucketName),
		Key:    aws.String(key),
		Body:   bytes.NewReader(p),
	})
	if err != nil {
		return fmt.Errorf("failed to upload file %s to s3: %w", path, err)
	}
	return nil
}

func (a *AWSS3) Delete(ctx context.Context, key string) error {
	_, err := a.cli.DeleteObject(ctx, &s3.DeleteObjectInput{
		Bucket: aws.String(a.bucketName),
		Key:    aws.String(key),
	})
	if err != nil {
		//ignore the error if the file is not found
		if strings.Contains(err.Error(), "NoSuchKey") {
			return nil
		}
		return fmt.Errorf("failed to delete file %s from s3: %w", key, err)
	}
	return nil
}

func (a *AWSS3) DownloadObjects(ctx context.Context, key string) ([]byte, error) {
	resp, err := a.cli.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(a.bucketName),
		Key:    aws.String(key),
	})
	if err != nil {
		return nil, fmt.Errorf("failed to download file %s from s3: %w", key, err)
	}
	return io.ReadAll(resp.Body)
}

type MockMeasureSender struct {
	t      *testing.T
	sended []models.BillingReport
}

func (m *MockMeasureSender) InitMeasureHandler(env string, vendor string, skipclusterid, usesqs bool, url ...string) error {
	assert.Equal(m.t, env, constants.EnvDev)
	assert.Equal(m.t, vendor, constants.VendorAWS)
	assert.Equal(m.t, skipclusterid, true)
	assert.Equal(m.t, usesqs, true)
	assert.Equal(m.t, url, []string{})
	return nil
}

func (m *MockMeasureSender) SendMeasure(ctx context.Context, billingMeasure models.BillingReport) error {
	m.sended = append(m.sended, billingMeasure)
	return nil
}

func (m *MockMeasureSender) GetSended() []models.BillingReport {
	return m.sended
}

func (m *MockMeasureSender) CheckMeasure(billingMeasure models.BillingReport) error {
	assert.True(m.t, billingMeasure.TenantID >= 1)
	assert.True(m.t, len(billingMeasure.CostID) >= 1)
	assert.True(m.t, billingMeasure.Cost >= 1)
	assert.True(m.t, billingMeasure.CostTimestamp >= 1)
	assert.True(m.t, len(billingMeasure.Unit) >= 1)
	return nil
}