// Copyright 2025 Microsoft Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package corecosmosstoragetesting

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	azcorearm "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
	"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"

	"github.com/Azure/ARO-HCP/internal/api/coreapi"
	"github.com/Azure/ARO-HCP/internal/api/metadataapi"
	"github.com/Azure/ARO-HCP/internal/database/cosmosstorage/corecosmosstorage"
	"github.com/Azure/ARO-HCP/internal/database/cosmosstorage/cosmosstorageutils"
)

// MockResourcesDBClient implements the corecosmosstorage.ResourcesDBClient interface for unit testing.
// It stores documents in memory and supports loading from cosmos-record context directories.
type MockResourcesDBClient struct {
	mu sync.RWMutex

	// documents stores all documents keyed by their cosmos ID
	documents map[string]json.RawMessage

	// globalListers is an optional custom global listers implementation for testing
	globalListers corecosmosstorage.ResourcesGlobalListers

	// changeFeed records each successful StoreDocument call so the
	// production change-feed consumer can be exercised against the
	// in-memory mock the same way it is against real Cosmos.
	changeFeed MockChangeFeed
}

// NewMockResourcesDBClient creates a new mock ResourcesDBClient with empty storage.
func NewMockResourcesDBClient() *MockResourcesDBClient {
	return &MockResourcesDBClient{
		documents: make(map[string]json.RawMessage),
	}
}

// SetResourcesGlobalListers sets a custom global listers implementation for testing.
// This allows tests to provide custom ResourcesGlobalListers that return errors or paginate.
func (m *MockResourcesDBClient) SetResourcesGlobalListers(globalListers corecosmosstorage.ResourcesGlobalListers) {
	m.globalListers = globalListers
}

// NewTransaction creates a new mock transaction.
func (m *MockResourcesDBClient) NewTransaction(pk string) cosmosstorageutils.DBTransaction {
	return newMockTransaction(pk, m)
}

// UntypedCRUD provides access to untyped resource operations.
func (m *MockResourcesDBClient) UntypedCRUD(parentResourceID azcorearm.ResourceID) (cosmosstorageutils.UntypedResourceCRUD, error) {
	return newMockUntypedCRUD(m, parentResourceID), nil
}

// HCPClusters returns a CRUD interface for HCPCluster resources.
func (m *MockResourcesDBClient) HCPClusters(subscriptionID, resourceGroupName string) corecosmosstorage.HCPClusterCRUD {
	var parentResourceID *azcorearm.ResourceID
	if len(resourceGroupName) > 0 {
		parentResourceID = metadataapi.Must(coreapi.ToResourceGroupResourceID(subscriptionID, resourceGroupName))
	} else {
		parentResourceID = metadataapi.Must(coreapi.ToSubscriptionResourceID(subscriptionID))
	}

	return newMockHCPClusterCRUD(m, parentResourceID)
}

// Operations returns a CRUD interface for operation resources.
func (m *MockResourcesDBClient) Operations(subscriptionID string) corecosmosstorage.OperationCRUD {
	parentResourceID := metadataapi.Must(coreapi.ToSubscriptionResourceID(subscriptionID))

	return newMockOperationCRUD(m, parentResourceID)
}

// Subscriptions returns a CRUD interface for subscription resources.
func (m *MockResourcesDBClient) Subscriptions() cosmosstorageutils.ResourceCRUD[coreapi.Subscription, *coreapi.Subscription] {
	return newMockSubscriptionCRUD(m)
}

// ListMissingResourceID returns documents that lack a resourceID field.
func (m *MockResourcesDBClient) ListMissingResourceID(ctx context.Context, options *cosmosstorageutils.DBClientListResourceDocsOptions) (cosmosstorageutils.DBClientIterator[cosmosstorageutils.TypedDocument], error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	var ids []string
	var items []*cosmosstorageutils.TypedDocument
	for _, data := range m.documents {
		var typedDoc cosmosstorageutils.TypedDocument
		if err := json.Unmarshal(data, &typedDoc); err != nil {
			continue
		}
		if typedDoc.ResourceID != nil {
			continue
		}
		ids = append(ids, typedDoc.ID)
		items = append(items, &typedDoc)
	}
	return NewMockIterator(ids, items), nil
}

// ResourcesGlobalListers returns interfaces for listing all resources of a particular
// type across all partitions. If a custom ResourcesGlobalListers was set via SetResourcesGlobalListers,
// that is returned instead.
func (m *MockResourcesDBClient) ResourcesGlobalListers() corecosmosstorage.ResourcesGlobalListers {
	if m.globalListers != nil {
		return m.globalListers
	}
	return &mockResourcesGlobalListers{client: m}
}

// ServiceProviderClusters returns a CRUD interface for service provider cluster resources.
func (m *MockResourcesDBClient) ServiceProviderClusters(subscriptionID, resourceGroupName, clusterName string) cosmosstorageutils.ResourceCRUD[coreapi.ServiceProviderCluster, *coreapi.ServiceProviderCluster] {
	clusterResourceID := metadataapi.Must(coreapi.ToClusterResourceID(subscriptionID, resourceGroupName, clusterName))
	return newMockServiceProviderClusterCRUD(m, clusterResourceID)
}

// ServiceProviderNodePools returns a CRUD interface for service provider node pool resources.
func (m *MockResourcesDBClient) ServiceProviderNodePools(subscriptionID, resourceGroupName, clusterName, nodePoolName string) cosmosstorageutils.ResourceCRUD[coreapi.ServiceProviderNodePool, *coreapi.ServiceProviderNodePool] {
	nodePoolResourceID := metadataapi.Must(coreapi.ToNodePoolResourceID(subscriptionID, resourceGroupName, clusterName, nodePoolName))
	return newMockServiceProviderNodePoolCRUD(m, nodePoolResourceID)
}

// ReadChangeFeed reads the in-memory change-feed log. Each
// successful StoreDocument call records a snapshot of the document;
// reads return everything past the position encoded in
// options.Continuation. Hard deletes are not recorded, which mirrors
// "latest version" mode in real Cosmos DB. Soft deletes go through
// StoreDocument and are therefore recorded.
func (m *MockResourcesDBClient) ReadChangeFeed(ctx context.Context, options *azcosmos.ChangeFeedOptions) (azcosmos.ChangeFeedResponse, error) {
	var continuation string
	if options != nil && options.Continuation != nil {
		continuation = *options.Continuation
	}
	items, nextToken, hasNew := m.changeFeed.Read(continuation)
	return BuildMockChangeFeedResponse(items, nextToken, hasNew), nil
}

// ReadFeedRanges returns the single feed range the mock
// advertises. Real Cosmos may report many ranges; one is enough for
// the in-memory mock because there is no partition-level parallelism
// to model.
func (m *MockResourcesDBClient) ReadFeedRanges(ctx context.Context, options *azcosmos.FeedRangesOptions) ([]azcosmos.FeedRange, error) {
	return []azcosmos.FeedRange{MockChangeFeedFeedRange}, nil
}

// LoadFromDirectory loads cosmos-record context data from a directory.
// It reads all JSON files that match the pattern for "load" directories.
func (m *MockResourcesDBClient) LoadFromDirectory(dirPath string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	return filepath.Walk(dirPath, func(filePath string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Skip directories
		if info.IsDir() {
			return nil
		}

		// Only process JSON files
		if !strings.HasSuffix(strings.ToLower(filePath), ".json") {
			return nil
		}

		// Read the file
		data, err := os.ReadFile(filePath)
		if err != nil {
			return fmt.Errorf("failed to read file %s: %w", filePath, err)
		}

		// Parse as TypedDocument to get the ID
		var typedDoc cosmosstorageutils.TypedDocument
		if err := json.Unmarshal(data, &typedDoc); err != nil {
			return fmt.Errorf("failed to unmarshal file %s: %w", filePath, err)
		}

		// Store the document
		if len(typedDoc.ID) != 0 {
			m.documents[strings.ToLower(typedDoc.ID)] = data
		}

		return nil
	})
}

// LoadContent loads a single JSON document into the mock database.
// This implements the ContentLoader interface from integrationutils.
func (m *MockResourcesDBClient) LoadContent(ctx context.Context, content []byte) error {
	// Parse as TypedDocument to get the ID
	var typedDoc cosmosstorageutils.TypedDocument
	if err := json.Unmarshal(content, &typedDoc); err != nil {
		return fmt.Errorf("failed to unmarshal content: %w", err)
	}

	if len(typedDoc.ID) == 0 {
		return fmt.Errorf("document is missing ID field")
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	m.documents[strings.ToLower(typedDoc.ID)] = content
	return nil
}

// ListAllDocuments returns all documents in the mock database.
// This implements the DocumentLister interface from integrationutils.
func (m *MockResourcesDBClient) ListAllDocuments(ctx context.Context) ([]*cosmosstorageutils.TypedDocument, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	var results []*cosmosstorageutils.TypedDocument
	for _, data := range m.documents {
		var typedDoc cosmosstorageutils.TypedDocument
		if err := json.Unmarshal(data, &typedDoc); err != nil {
			return nil, fmt.Errorf("failed to unmarshal document: %w", err)
		}
		results = append(results, &typedDoc)
	}
	return results, nil
}

// StoreDocument stores a raw JSON document in the mock database and
// appends it to the in-memory change-feed log so that consumers of
// ReadChangeFeed see the mutation.
func (m *MockResourcesDBClient) StoreDocument(cosmosID string, data json.RawMessage) {
	m.mu.Lock()
	m.documents[strings.ToLower(cosmosID)] = data
	m.mu.Unlock()
	m.changeFeed.Record(data)
}

// GetDocument retrieves a raw JSON document from the mock database.
func (m *MockResourcesDBClient) GetDocument(cosmosID string) (json.RawMessage, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	data, ok := m.documents[strings.ToLower(cosmosID)]
	return data, ok
}

// DeleteDocument removes a document from the mock database.
func (m *MockResourcesDBClient) DeleteDocument(cosmosID string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.documents, strings.ToLower(cosmosID))
}

// ListDocuments returns all documents matching the given resource type and prefix.
// Soft-deleted documents (those with a DeletionTimestamp set) are excluded.
func (m *MockResourcesDBClient) ListDocuments(resourceType *azcorearm.ResourceType, prefix string) []json.RawMessage {
	m.mu.RLock()
	defer m.mu.RUnlock()

	var results []json.RawMessage
	for _, data := range m.documents {
		var typedDoc cosmosstorageutils.TypedDocument
		if err := json.Unmarshal(data, &typedDoc); err != nil {
			continue
		}

		// Mirror the production query, which requires IS_DEFINED(c.resourceID);
		// documents without a resourceID are never returned by list.
		if typedDoc.ResourceID == nil {
			continue
		}

		if typedDoc.DeletionTimestamp != nil {
			continue
		}

		// Check resource type match if specified
		if resourceType != nil {
			if !strings.EqualFold(typedDoc.ResourceType, resourceType.String()) {
				continue
			}
		}

		// Check prefix match if specified.  /subscriptions/ doesn't count because in real storage we pass no prefix for subscriptions
		if len(prefix) != 0 && prefix != "/subscriptions/" {
			if !strings.HasPrefix(strings.ToLower(typedDoc.ResourceID.String()), strings.ToLower(prefix)) {
				continue
			}
		}

		results = append(results, data)
	}

	return results
}

// Clear removes all documents from the mock database.
func (m *MockResourcesDBClient) Clear() {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.documents = make(map[string]json.RawMessage)
}

// GetAllDocuments returns a copy of all documents (for testing purposes).
func (m *MockResourcesDBClient) GetAllDocuments() map[string]json.RawMessage {
	m.mu.RLock()
	defer m.mu.RUnlock()

	result := make(map[string]json.RawMessage, len(m.documents))
	for k, v := range m.documents {
		result[k] = v
	}
	return result
}

var _ corecosmosstorage.ResourcesDBClient = &MockResourcesDBClient{}

// mockTransaction implements cosmosstorageutils.DBTransaction for the mock client.
type mockTransaction struct {
	pk        string
	client    *MockResourcesDBClient
	steps     []mockTransactionStep
	onSuccess []cosmosstorageutils.DBTransactionCallback
}

type mockTransactionStep struct {
	details cosmosstorageutils.CosmosDBTransactionStepDetails
	execute func() (string, json.RawMessage, error)
}

func newMockTransaction(pk string, client *MockResourcesDBClient) *mockTransaction {
	return &mockTransaction{
		pk:     strings.ToLower(pk),
		client: client,
	}
}

func (t *mockTransaction) GetPartitionKey() string {
	return t.pk
}

func (t *mockTransaction) AddStep(details cosmosstorageutils.CosmosDBTransactionStepDetails, stepFn cosmosstorageutils.CosmosDBTransactionStep) {
	// We need to capture what the step does for the mock
	t.steps = append(t.steps, mockTransactionStep{
		details: details,
		execute: func() (string, json.RawMessage, error) {
			// The real implementation uses TransactionalBatch, but we just execute directly
			// We'll handle this in Execute by storing the details
			return details.CosmosID, nil, nil
		},
	})
}

func (t *mockTransaction) OnSuccess(callback cosmosstorageutils.DBTransactionCallback) {
	if callback != nil {
		t.onSuccess = append(t.onSuccess, callback)
	}
}

func (t *mockTransaction) Execute(ctx context.Context, o *azcosmos.TransactionalBatchOptions) (cosmosstorageutils.DBTransactionResult, error) {
	result := &mockTransactionResult{
		items: make(map[string]json.RawMessage),
	}

	// Execute all steps
	for step, s := range t.steps {
		cosmosID, data, err := s.execute()
		if err != nil {
			var responseErr *azcore.ResponseError
			if errors.As(err, &responseErr) {
				return nil, cosmosstorageutils.NewTransactionStepError(step+1, len(t.steps), responseErr.StatusCode, s.details)
			}
			return nil, err
		}
		if data != nil {
			result.items[cosmosID] = data
		}
	}

	// Call success callbacks
	for _, callback := range t.onSuccess {
		callback(result)
	}

	return result, nil
}

var _ cosmosstorageutils.DBTransaction = &mockTransaction{}

// mockTransactionResult implements cosmosstorageutils.DBTransactionResult.
type mockTransactionResult struct {
	items map[string]json.RawMessage
}

func (r *mockTransactionResult) GetItem(cosmosUID string) (any, error) {
	data, ok := r.items[cosmosUID]
	if !ok {
		return nil, corecosmosstorage.ErrItemNotFound
	}

	var typedDoc cosmosstorageutils.TypedDocument
	if err := json.Unmarshal(data, &typedDoc); err != nil {
		return nil, err
	}

	switch strings.ToLower(typedDoc.ResourceType) {
	case strings.ToLower(coreapi.ClusterResourceType.String()):
		var cosmosObj cosmosstorageutils.GenericDocument[coreapi.HCPOpenShiftCluster]
		if err := json.Unmarshal(data, &cosmosObj); err != nil {
			return nil, err
		}
		return cosmosstorageutils.CosmosGenericToInternal(&cosmosObj)
	case strings.ToLower(coreapi.NodePoolResourceType.String()):
		var cosmosObj cosmosstorageutils.GenericDocument[coreapi.HCPOpenShiftClusterNodePool]
		if err := json.Unmarshal(data, &cosmosObj); err != nil {
			return nil, err
		}
		return cosmosstorageutils.CosmosGenericToInternal(&cosmosObj)
	case strings.ToLower(coreapi.ExternalAuthResourceType.String()):
		var cosmosObj cosmosstorageutils.GenericDocument[coreapi.HCPOpenShiftClusterExternalAuth]
		if err := json.Unmarshal(data, &cosmosObj); err != nil {
			return nil, err
		}
		return cosmosstorageutils.CosmosGenericToInternal(&cosmosObj)
	default:
		return nil, fmt.Errorf("unknown resource type '%s'", typedDoc.ResourceType)
	}
}

var _ cosmosstorageutils.DBTransactionResult = &mockTransactionResult{}

// MockIterator implements cosmosstorageutils.DBClientIterator for in-memory iteration.
type MockIterator[T any] struct {
	items             []*T
	ids               []string
	continuationToken string
	err               error
	index             int
}

func NewMockIterator[T any](ids []string, items []*T) *MockIterator[T] {
	return &MockIterator[T]{
		items: items,
		ids:   ids,
		index: 0,
	}
}

func (iter *MockIterator[T]) Items(ctx context.Context) cosmosstorageutils.DBClientIteratorItem[T] {
	return func(yield func(string, *T) bool) {
		for i, item := range iter.items {
			if !yield(iter.ids[i], item) {
				return
			}
		}
	}
}

func (iter *MockIterator[T]) GetContinuationToken() string {
	return iter.continuationToken
}

func (iter *MockIterator[T]) GetError() error {
	return iter.err
}

var _ cosmosstorageutils.DBClientIterator[coreapi.HCPOpenShiftCluster] = &MockIterator[coreapi.HCPOpenShiftCluster]{}
