package nextgenrepoter

import (
	"context"
	"encoding/json"
	"errors"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/metering/pkg/common/config"
	"github.com/tidbcloud/metering/pkg/common/models"
	"go.uber.org/zap"
	"go.uber.org/zap/zaptest/observer"
)

func TestNewNextGenBillingRepoter(t *testing.T) {
	logger := zap.NewNop()
	conf := &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{}, Vendor: "aws"}

	originalLoader := newNextGenRepoterLoaderFunc
	originalSenderFactory := measureSenderFactory
	t.Cleanup(func() {
		newNextGenRepoterLoaderFunc = originalLoader
		measureSenderFactory = originalSenderFactory
	})

	mockLoader := &fakeLoader{}
	newNextGenRepoterLoaderFunc = func(conf *config.ServiceConfig, log *zap.Logger) (NextGenRepoterLoader, error) {
		return mockLoader, nil
	}

	mockSender := &fakeSender{}
	measureSenderFactory = func() measureSender { return mockSender }

	reporter, err := NewNextGenBillingRepoter(conf, logger)
	require.NoError(t, err)
	assert.Equal(t, mockLoader, reporter.loader)
	assert.Equal(t, mockSender, reporter.sender)

	measureSenderFactory = func() measureSender { return &fakeSender{initErr: errors.New("init failed")} }
	_, err = NewNextGenBillingRepoter(conf, logger)
	require.Error(t, err)

	newNextGenRepoterLoaderFunc = func(conf *config.ServiceConfig, log *zap.Logger) (NextGenRepoterLoader, error) {
		return nil, errors.New("loader failed")
	}
	measureSenderFactory = originalSenderFactory
	_, err = NewNextGenBillingRepoter(conf, logger)
	require.Error(t, err)

	_, err = NewNextGenBillingRepoter(&config.ServiceConfig{}, logger)
	require.Error(t, err)
}

func TestNextGenBillingRepoter_Report(t *testing.T) {
	now := time.Now().Unix()
	logger := zap.NewNop()
	conf := &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{}, Vendor: "aws"}

	t.Run("checkpoint_error", func(t *testing.T) {
		loader := &fakeLoader{loadCheckpointErr: errors.New("load fail")}
		reporter := &NextGenBillingRepoter{conf: conf, loader: loader, sender: &fakeSender{}, log: logger}
		err := reporter.Report(context.Background())
		require.Error(t, err)
	})

	t.Run("checkpoint_up_to_date", func(t *testing.T) {
		loader := &fakeLoader{checkpoint: now + 1000}
		reporter := &NextGenBillingRepoter{conf: conf, loader: loader, sender: &fakeSender{}, log: logger}
		err := reporter.Report(context.Background())
		require.NoError(t, err)
		assert.False(t, loader.loadCalled)
	})

	t.Run("load_error", func(t *testing.T) {
		loader := &fakeLoader{checkpoint: now - 10, loadErr: errors.New("load fail")}
		reporter := &NextGenBillingRepoter{conf: conf, loader: loader, sender: &fakeSender{}, log: logger}
		err := reporter.Report(context.Background())
		require.Error(t, err)
		assert.True(t, loader.loadCalled)
	})

	t.Run("no_data", func(t *testing.T) {
		loader := &fakeLoader{checkpoint: now - 10}
		reporter := &NextGenBillingRepoter{conf: conf, loader: loader, sender: &fakeSender{}, log: logger}
		err := reporter.Report(context.Background())
		require.NoError(t, err)
		assert.True(t, loader.loadCalled)
		assert.Equal(t, int64(0), loader.savedCheckpoint)
	})

	t.Run("convert_error", func(t *testing.T) {
		loader := &fakeLoader{
			checkpoint: now - 10,
			reports: []*models.BillingReportV2Deltalake{
				{TenantID: 1, CostID: "c1", CostTimestamp: now - 5, CreatedAt: now - 5, Unit: "u", UnitValue: "bad", Cost: 1},
			},
		}
		reporter := &NextGenBillingRepoter{conf: conf, loader: loader, sender: &fakeSender{}, log: logger}
		err := reporter.Report(context.Background())
		require.Error(t, err)
		assert.True(t, loader.loadCalled)
		assert.Equal(t, int64(0), loader.savedCheckpoint)
	})

	t.Run("send_error", func(t *testing.T) {
		report := &models.BillingReportV2Deltalake{TenantID: 1, CostID: "c1", CostTimestamp: now - 5, CreatedAt: now - 5, Unit: "u", UnitValue: "1", Cost: 2}
		loader := &fakeLoader{checkpoint: now - 10, reports: []*models.BillingReportV2Deltalake{report}}
		sender := &fakeSender{sendErr: errors.New("send fail")}
		reporter := &NextGenBillingRepoter{conf: conf, loader: loader, sender: sender, log: logger}
		err := reporter.Report(context.Background())
		require.Error(t, err)
		assert.Equal(t, 1, sender.sendCount)
		assert.Equal(t, int64(0), loader.savedCheckpoint)
	})

	t.Run("success_updates_checkpoint", func(t *testing.T) {
		report1 := &models.BillingReportV2Deltalake{TenantID: 1, CostID: "c1", CostTimestamp: now - 5, CreatedAt: now - 5, Unit: "u", UnitValue: "1", Cost: 2}
		report2 := &models.BillingReportV2Deltalake{TenantID: 1, CostID: "c2", CostTimestamp: now - 3, CreatedAt: now - 2, Unit: "u", UnitValue: "2", Cost: 2}
		loader := &fakeLoader{checkpoint: now - 10, reports: []*models.BillingReportV2Deltalake{report1, report2}}
		sender := &fakeSender{}
		reporter := &NextGenBillingRepoter{conf: conf, loader: loader, sender: sender, log: logger}
		err := reporter.Report(context.Background())
		require.NoError(t, err)
		assert.Equal(t, 2, sender.sendCount)
		assert.Equal(t, report2.CreatedAt, loader.savedCheckpoint)
	})

	t.Run("logs_billing_record_for_analysis_before_send", func(t *testing.T) {
		core, observedLogs := observer.New(zap.DebugLevel)
		logger := zap.New(core)
		report := &models.BillingReportV2Deltalake{
			TenantID:      123,
			CostID:        "c1",
			CostTimestamp: now - 5,
			CreatedAt:     now - 5,
			Unit:          "u",
			UnitValue:     "2",
			Cost:          3,
			Labels:        `{"cluster_id":"cluster-1"}`,
		}
		loader := &fakeLoader{checkpoint: now - 10, reports: []*models.BillingReportV2Deltalake{report}}
		sender := &fakeSender{}
		reporter := &NextGenBillingRepoter{conf: conf, loader: loader, sender: sender, log: logger}

		err := reporter.Report(context.Background())
		require.NoError(t, err)

		entries := observedLogs.FilterMessage("billing records for analysis").AllUntimed()
		require.Len(t, entries, 1)

		expectedBillingReport, err := report.Convert2BillingReport()
		require.NoError(t, err)

		fields := entries[0].ContextMap()
		recordJSON, err := json.Marshal(fields["billing_record"])
		require.NoError(t, err)

		expectedJSON, err := json.Marshal(expectedBillingReport)
		require.NoError(t, err)

		assert.JSONEq(t, string(expectedJSON), string(recordJSON))
	})
}

func TestGroupReportsByHour(t *testing.T) {
	logger := zap.NewNop()
	reporter := &NextGenBillingRepoter{log: logger}

	hour1 := time.Date(2024, 1, 1, 10, 15, 0, 0, time.UTC)
	hour2 := hour1.Add(90 * time.Minute)

	reports := []*models.BillingReportV2Deltalake{
		{CostTimestamp: hour1.Unix()},
		{CostTimestamp: hour1.Add(30 * time.Minute).Unix()},
		{CostTimestamp: hour2.Unix()},
	}

	batches, err := reporter.groupReportsByHour(reports)
	require.NoError(t, err)
	require.Len(t, batches, 2)
	assert.Equal(t, hour1.Truncate(time.Hour).Unix(), batches[0].StartTime)
	assert.Len(t, batches[0].Reports, 2)
	assert.Len(t, batches[1].Reports, 1)
	assert.Less(t, batches[0].StartTime, batches[1].StartTime)
}

func TestSendBatch(t *testing.T) {
	ctx := context.Background()
	logger := zap.NewNop()
	sender := &fakeSender{}
	batch := &HourlyBatch{
		Hour: time.Unix(1700000000, 0),
		Reports: []*models.BillingReportV2Deltalake{
			{TenantID: 1, CostID: "c1", CostTimestamp: 1, CreatedAt: 1, Unit: "u", UnitValue: "1", Cost: 2},
			{TenantID: 1, CostID: "c2", CostTimestamp: 2, CreatedAt: 2, Unit: "u", UnitValue: "3", Cost: 4},
		},
	}
	reporter := &NextGenBillingRepoter{sender: sender, log: logger}

	err := reporter.sendBatch(ctx, batch)
	require.NoError(t, err)
	assert.Equal(t, 2, sender.sendCount)

	sender.sendErr = errors.New("send")
	err = reporter.sendBatch(ctx, batch)
	require.Error(t, err)
	assert.GreaterOrEqual(t, sender.sendCount, 2)
}

type fakeLoader struct {
	checkpoint        int64
	loadErr           error
	loadCheckpointErr error
	saveErr           error
	reports           []*models.BillingReportV2Deltalake
	loadCalled        bool
	savedCheckpoint   int64
}

func (f *fakeLoader) LoadMeteringDataInTimeRange(ctx context.Context, startTime, endTime int64) ([]*models.BillingReportV2Deltalake, error) {
	f.loadCalled = true
	if f.loadErr != nil {
		return nil, f.loadErr
	}
	return f.reports, nil
}

func (f *fakeLoader) LoadMeteringDataInTimeRangeWithDeltaRS(ctx context.Context, startTime, endTime int64) ([]*models.BillingReportV2Deltalake, error) {
	// For testing purposes, use the same logic as the DuckDB method
	return f.LoadMeteringDataInTimeRange(ctx, startTime, endTime)
}

func (f *fakeLoader) LoadCheckpoint(ctx context.Context) (int64, error) {
	if f.loadCheckpointErr != nil {
		return 0, f.loadCheckpointErr
	}
	return f.checkpoint, nil
}

func (f *fakeLoader) SaveCheckpoint(ctx context.Context, checkpoint int64) error {
	if f.saveErr != nil {
		return f.saveErr
	}
	f.savedCheckpoint = checkpoint
	return nil
}

type fakeSender struct {
	initErr   error
	sendErr   error
	sendCount int
}

func (f *fakeSender) InitMeasureHandler(env string, vendor string, skipclusterid, usesqs bool, url ...string) error {
	return f.initErr
}

func (f *fakeSender) SendMeasure(ctx context.Context, billingMeasure models.BillingReport) error {
	f.sendCount++
	if f.sendErr != nil {
		return f.sendErr
	}
	return nil
}

func (f *fakeSender) CheckMeasure(billingMeasure models.BillingReport) error {
	return nil
}

var _ measureSender = (*fakeSender)(nil)
var _ NextGenRepoterLoader = (*fakeLoader)(nil)