package service

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"

	coremodel "security-center-backend/internal/model"
	"security-center-backend/internal/modules/caasm/mocks"
	caasmmodel "security-center-backend/internal/modules/caasm/model"
	"security-center-backend/internal/modules/caasm/repository"
)

type countingAssetHistoryRepo struct {
	calls   int
	entries []*caasmmodel.AssetHistory
}

func (r *countingAssetHistoryRepo) Create(ctx context.Context, entry *caasmmodel.AssetHistory) error {
	r.calls++
	if entry != nil {
		copyEntry := *entry
		r.entries = append(r.entries, &copyEntry)
	}
	return nil
}

type countingTagRepo struct {
	calls   int
	assetID uint64
	tags    map[string]string
}

func (r *countingTagRepo) ReplaceTags(ctx context.Context, assetID uint64, tags map[string]string) error {
	r.calls++
	r.assetID = assetID
	r.tags = make(map[string]string, len(tags))
	for k, v := range tags {
		r.tags[k] = v
	}
	return nil
}

func (r *countingTagRepo) GetTags(ctx context.Context, assetID uint64) (map[string]string, error) {
	return nil, nil
}

type sequencedDNSRecordRepo struct {
	existing      []caasmmodel.DNSRecord
	replaceErrs   []error
	replaceCalls  int
	replacedZone  uint64
	replacedSlice []caasmmodel.DNSRecord
}

func (r *sequencedDNSRecordRepo) List(ctx context.Context, filter repository.DNSRecordFilter, page, pageSize int) ([]caasmmodel.DNSRecord, int64, error) {
	return nil, 0, nil
}

func (r *sequencedDNSRecordRepo) GetByDomain(ctx context.Context, domainID uint64) ([]caasmmodel.DNSRecord, error) {
	records := make([]caasmmodel.DNSRecord, len(r.existing))
	copy(records, r.existing)
	for i := range records {
		records[i].DomainAssetID = domainID
	}
	return records, nil
}

func (r *sequencedDNSRecordRepo) ReplaceForDomain(ctx context.Context, domainID uint64, records []caasmmodel.DNSRecord) error {
	r.replaceCalls++
	r.replacedZone = domainID
	r.replacedSlice = append([]caasmmodel.DNSRecord(nil), records...)
	if r.replaceCalls <= len(r.replaceErrs) {
		return r.replaceErrs[r.replaceCalls-1]
	}
	return nil
}

func (r *sequencedDNSRecordRepo) UpdateOwner(ctx context.Context, id uint64, ownerEmail, deptID string, employeeID *uint64, ignoreChanges *bool, suppressPattern *string) error {
	return nil
}

func TestPersistSyncPayloadDoesNotRetryAwayExecutionCancellation(t *testing.T) {
	t.Run("user_stop", func(t *testing.T) {
		ctx, cancel := context.WithCancel(context.Background())
		cancel()

		mockJobRepo := new(mocks.MockSyncJobRepository)
		mockAssetRepo := new(mocks.MockAssetRepository)
		svc := NewCloudAccountService(nil, mockJobRepo, mockAssetRepo)

		job := &caasmmodel.SyncJob{Base: coremodel.Base{ID: 501}, Metadata: coremodel.JSONB{}}
		payload := &caasmmodel.IngestPayload{
			CloudAccountID: 99,
			Assets: []caasmmodel.Asset{{
				NativeID:  "stopped-asset",
				Name:      "stopped-asset",
				AssetType: caasmmodel.AssetTypeNetwork,
				Service:   "EC2",
				Status:    caasmmodel.AssetStatusActive,
			}},
		}

		mockJobRepo.On("UpdateFields", mock.Anything, uint64(501), mock.Anything).Return(nil).Once()

		err := svc.persistSyncPayload(ctx, job, payload)
		require.Error(t, err)
		require.True(t, errors.Is(err, context.Canceled))
		mockAssetRepo.AssertNotCalled(t, "FindMinimalByNativeIDs", mock.Anything, uint64(99), []string{"stopped-asset"})
		mockJobRepo.AssertExpectations(t)
	})

	t.Run("execution_timeout", func(t *testing.T) {
		ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second))
		defer cancel()

		mockJobRepo := new(mocks.MockSyncJobRepository)
		mockAssetRepo := new(mocks.MockAssetRepository)
		svc := NewCloudAccountService(nil, mockJobRepo, mockAssetRepo)

		job := &caasmmodel.SyncJob{Base: coremodel.Base{ID: 502}, Metadata: coremodel.JSONB{}}
		payload := &caasmmodel.IngestPayload{
			CloudAccountID: 100,
			Assets: []caasmmodel.Asset{{
				NativeID:  "timed-out-asset",
				Name:      "timed-out-asset",
				AssetType: caasmmodel.AssetTypeNetwork,
				Service:   "EC2",
				Status:    caasmmodel.AssetStatusActive,
			}},
		}

		mockJobRepo.On("UpdateFields", mock.Anything, uint64(502), mock.Anything).Return(nil).Once()

		err := svc.persistSyncPayload(ctx, job, payload)
		require.Error(t, err)
		require.True(t, errors.Is(err, context.DeadlineExceeded))
		mockAssetRepo.AssertNotCalled(t, "FindMinimalByNativeIDs", mock.Anything, uint64(100), []string{"timed-out-asset"})
		mockJobRepo.AssertExpectations(t)
	})
}

func TestPersistSyncPayloadRetryDoesNotDuplicateHistoryTagsOrProgress(t *testing.T) {
	ctx := context.Background()

	t.Run("asset_history_and_tags", func(t *testing.T) {
		mockJobRepo := new(mocks.MockSyncJobRepository)
		mockAssetRepo := new(mocks.MockAssetRepository)
		historyRepo := &countingAssetHistoryRepo{}
		tagRepo := &countingTagRepo{}
		svc := NewCloudAccountService(nil, mockJobRepo, mockAssetRepo)
		svc.SetHistoryRepository(historyRepo)
		svc.SetTagRepository(tagRepo)

		job := &caasmmodel.SyncJob{Base: coremodel.Base{ID: 601}, Metadata: coremodel.JSONB{}}
		accountID := uint64(201)
		payload := &caasmmodel.IngestPayload{
			CloudAccountID: accountID,
			Assets: []caasmmodel.Asset{{
				CloudAccountID: &accountID,
				NativeID:       "asset-retry-1",
				Name:           "asset-retry-1",
				AssetType:      caasmmodel.AssetTypeNetwork,
				Service:        "EC2",
				Status:         caasmmodel.AssetStatusActive,
				Metadata: coremodel.JSONB{
					"tags": map[string]string{
						"env":   "prod",
						"owner": "secops",
					},
				},
			}},
		}

		mockJobRepo.On("UpdateFields", mock.Anything, uint64(601), mock.Anything).Return(nil).Once()
		mockAssetRepo.On("FindMinimalByNativeIDs", mock.Anything, accountID, []string{"asset-retry-1"}).Return([]repository.AssetMinimal{}, nil).Once()
		mockAssetRepo.On("SaveAssetSnapshot", mock.Anything, mock.MatchedBy(func(asset *caasmmodel.Asset) bool {
			return asset.NativeID == "asset-retry-1" && asset.Metadata != nil
		})).Run(func(args mock.Arguments) {
			asset := args.Get(1).(*caasmmodel.Asset)
			asset.ID = 77
		}).Return(context.Canceled).Once()
		mockAssetRepo.On("SaveAssetSnapshot", mock.Anything, mock.MatchedBy(func(asset *caasmmodel.Asset) bool {
			return asset.NativeID == "asset-retry-1" && asset.Metadata != nil
		})).Run(func(args mock.Arguments) {
			asset := args.Get(1).(*caasmmodel.Asset)
			asset.ID = 77
		}).Return(nil).Once()

		require.NoError(t, svc.persistSyncPayload(ctx, job, payload))
		mockJobRepo.AssertNumberOfCalls(t, "UpdateFields", 1)
		assert.Equal(t, 2, historyRepo.calls)
		require.Len(t, historyRepo.entries, 2)
		assert.Equal(t, "upsert", historyRepo.entries[0].Action)
		assert.Equal(t, "tags.update", historyRepo.entries[1].Action)
		assert.Equal(t, 1, tagRepo.calls)
		assert.Equal(t, uint64(77), tagRepo.assetID)
		assert.Equal(t, "prod", tagRepo.tags["env"])
		assert.Equal(t, "secops", tagRepo.tags["owner"])
		mockAssetRepo.AssertExpectations(t)
		mockJobRepo.AssertExpectations(t)
	})

	t.Run("dns_event_emission", func(t *testing.T) {
		mockJobRepo := new(mocks.MockSyncJobRepository)
		mockAssetRepo := new(mocks.MockAssetRepository)
		dnsRepo := &sequencedDNSRecordRepo{
			existing: []caasmmodel.DNSRecord{{
				Base:          coremodel.Base{ID: 31},
				DomainAssetID: 9001,
				Name:          "www.example.com",
				Type:          "A",
				Values:        coremodel.StringArray{"10.0.0.1"},
				TTL:           300,
				OwnerEmail:    "ops@example.com",
				IgnoreChanges: false,
			}},
			replaceErrs: []error{context.Canceled},
		}
		eventRepo := &stubSyncJobEventRepo{}
		svc := NewCloudAccountService(nil, mockJobRepo, mockAssetRepo)
		svc.SetDNSRecordRepository(dnsRepo)
		svc.events = eventRepo

		job := &caasmmodel.SyncJob{Base: coremodel.Base{ID: 602}, Metadata: coremodel.JSONB{}}
		payload := &caasmmodel.IngestPayload{
			CloudAccountID: 202,
			DomainAssets: []caasmmodel.DomainAsset{{
				AssetID:     9001,
				DomainName:  "example.com",
				DomainClass: "hosted_zone",
				Registrar:   "route53",
				DNSRecords: coremodel.JSONB{
					"records": []map[string]interface{}{
						{
							"name":   "www.example.com",
							"type":   "A",
							"values": []string{"10.0.0.2"},
						},
					},
				},
			}},
		}

		mockJobRepo.On("UpdateFields", mock.Anything, uint64(602), mock.Anything).Return(nil).Once()
		mockAssetRepo.On("SaveDomainAssetSnapshot", mock.Anything, mock.MatchedBy(func(domainAsset *caasmmodel.DomainAsset) bool {
			return domainAsset.DomainName == "example.com" && domainAsset.AssetID == 9001
		})).Run(func(args mock.Arguments) {
			domainAsset := args.Get(1).(*caasmmodel.DomainAsset)
			domainAsset.ID = 9001
		}).Return(nil).Once()

		require.NoError(t, svc.persistSyncPayload(ctx, job, payload))
		assert.Len(t, eventRepo.events, 1)
		assert.Equal(t, "DNS target changed", eventRepo.events[0].Message)
		assert.Equal(t, 2, dnsRepo.replaceCalls)
		assert.Equal(t, uint64(9001), dnsRepo.replacedZone)
		mockAssetRepo.AssertExpectations(t)
		mockJobRepo.AssertExpectations(t)
	})
}

func TestPersistSyncPayloadSkipsDNSReplaceWhenRecordsMissing(t *testing.T) {
	ctx := context.Background()

	cases := []struct {
		name         string
		dnsRecords   coremodel.JSONB
		wantReplace  int
		wantReplaced int
	}{
		{name: "nil dns records means not collected"},
		{name: "missing records key means not collected", dnsRecords: coremodel.JSONB{"provider": "route53"}},
		{name: "nil records key means not collected", dnsRecords: coremodel.JSONB{"records": nil}},
		{name: "object records key means not collected", dnsRecords: coremodel.JSONB{"records": map[string]interface{}{"unexpected": "shape"}}},
		{name: "string records key means not collected", dnsRecords: coremodel.JSONB{"records": "collection failed"}},
		{name: "empty records key means explicit empty inventory", dnsRecords: coremodel.JSONB{"records": []interface{}{}}, wantReplace: 1},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			mockJobRepo := new(mocks.MockSyncJobRepository)
			mockAssetRepo := new(mocks.MockAssetRepository)
			dnsRepo := &sequencedDNSRecordRepo{
				existing: []caasmmodel.DNSRecord{{
					Base:          coremodel.Base{ID: 31},
					DomainAssetID: 9001,
					Name:          "www.example.com",
					Type:          "A",
					Values:        coremodel.StringArray{"10.0.0.1"},
				}},
			}
			svc := NewCloudAccountService(nil, mockJobRepo, mockAssetRepo)
			svc.SetDNSRecordRepository(dnsRepo)

			job := &caasmmodel.SyncJob{Base: coremodel.Base{ID: 603}, Metadata: coremodel.JSONB{}}
			payload := &caasmmodel.IngestPayload{
				CloudAccountID: 202,
				DomainAssets: []caasmmodel.DomainAsset{{
					AssetID:     9001,
					DomainName:  "example.com",
					DomainClass: "hosted_zone",
					Registrar:   "route53",
					DNSRecords:  tc.dnsRecords,
				}},
			}

			mockJobRepo.On("UpdateFields", mock.Anything, uint64(603), mock.Anything).Return(nil).Once()
			mockAssetRepo.On("SaveDomainAssetSnapshot", mock.Anything, mock.MatchedBy(func(domainAsset *caasmmodel.DomainAsset) bool {
				return domainAsset.DomainName == "example.com" && domainAsset.AssetID == 9001
			})).Run(func(args mock.Arguments) {
				domainAsset := args.Get(1).(*caasmmodel.DomainAsset)
				domainAsset.ID = 9001
			}).Return(nil).Once()

			require.NoError(t, svc.persistSyncPayload(ctx, job, payload))
			assert.Equal(t, tc.wantReplace, dnsRepo.replaceCalls)
			assert.Len(t, dnsRepo.replacedSlice, tc.wantReplaced)
			mockAssetRepo.AssertExpectations(t)
			mockJobRepo.AssertExpectations(t)
		})
	}
}

func TestQueueRetryClassificationForFailedSyncJobs(t *testing.T) {
	cases := []struct {
		name      string
		job       *caasmmodel.SyncJob
		wantRetry bool
	}{
		{
			name: "persist error is retryable",
			job: &caasmmodel.SyncJob{
				Status:       caasmmodel.SyncJobFailed,
				ErrorMessage: "persist failed",
				Metadata:     coremodel.JSONB{"termination_reason": "persist_error"},
			},
			wantRetry: true,
		},
		{
			name: "execution timeout is retryable",
			job: &caasmmodel.SyncJob{
				Status:       caasmmodel.SyncJobFailed,
				ErrorMessage: "timed out",
				Metadata:     coremodel.JSONB{"termination_reason": "execution_timeout"},
			},
			wantRetry: true,
		},
		{
			name: "partial terminal state is acked",
			job: &caasmmodel.SyncJob{
				Status:       caasmmodel.SyncJobPartial,
				ErrorMessage: "partial",
				Metadata:     coremodel.JSONB{"termination_reason": "partial_failures"},
			},
		},
		{
			name: "user stop is acked",
			job: &caasmmodel.SyncJob{
				Status:       caasmmodel.SyncJobFailed,
				ErrorMessage: "Job stopped by user",
				Metadata: coremodel.JSONB{
					"termination_reason":  "job_stopped",
					"cancellation_source": "user_stop",
				},
			},
		},
		{
			name: "account overlap is acked",
			job: &caasmmodel.SyncJob{
				Status:       caasmmodel.SyncJobFailed,
				ErrorMessage: "another sync is already running for this account",
				Metadata:     coremodel.JSONB{"coordination_state": "account_overlap"},
			},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			assert.Equal(t, tc.wantRetry, isQueueRetryableFailedSyncJob(tc.job))
		})
	}
}