// Copyright (c) Microsoft Corporation.
// Licensed under the Apache License 2.0.

package corecosmosstorage

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"

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

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

// ErrItemNotFound occurs when the requested item ID was not found,
// such as in a DBTransactionResult.
var ErrItemNotFound = errors.New("item not found")

// ErrWrongPartition occurs in a DBTransaction create step when the
// document has a partition key that differs from the transaction's
// partition key.
var ErrWrongPartition = errors.New("wrong partition key for transaction")

var _ cosmosstorageutils.DBTransaction = &cosmosDBTransaction{}

type cosmosDBTransaction struct {
	pk           string
	client       *azcosmos.ContainerClient
	steps        []cosmosstorageutils.CosmosDBTransactionStep
	stepsDetails []cosmosstorageutils.CosmosDBTransactionStepDetails
	onSuccess    []cosmosstorageutils.DBTransactionCallback
}

func newCosmosDBTransaction(pk string, client *azcosmos.ContainerClient) *cosmosDBTransaction {
	return &cosmosDBTransaction{
		pk:        strings.ToLower(pk),
		client:    client,
		steps:     nil,
		onSuccess: nil}
}

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

func (t *cosmosDBTransaction) AddStep(details cosmosstorageutils.CosmosDBTransactionStepDetails, stepFn cosmosstorageutils.CosmosDBTransactionStep) {
	t.steps = append(t.steps, stepFn)
	t.stepsDetails = append(t.stepsDetails, details)
}

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

func (t *cosmosDBTransaction) Execute(ctx context.Context, o *azcosmos.TransactionalBatchOptions) (cosmosstorageutils.DBTransactionResult, error) {
	logger := utils.LoggerFromContext(ctx)
	logger.Info("Executing transaction", "transaction", t.ToJSONStructRendering())

	result := newCosmosDBTransactionResult()

	if len(t.steps) > 0 {
		batch := t.client.NewTransactionalBatch(azcosmos.NewPartitionKeyString(t.pk))

		// Execute the queued steps to prepare the transaction. Collect
		// the item ID of each step to pair with the operation results.
		itemIDs := make([]string, 0, len(t.steps))
		for _, step := range t.steps {
			id, err := step(&batch)
			if err != nil {
				return nil, err
			}
			itemIDs = append(itemIDs, id)
		}

		response, err := t.client.ExecuteTransactionalBatch(ctx, batch, o)
		if err != nil {
			return nil, utils.TrackError(err)
		}

		if !response.Success {
			for step, result := range response.OperationResults {
				if result.StatusCode != http.StatusFailedDependency {
					logger.Error(nil, "transaction step failed",
						"step", step+1,
						"totalSteps", len(response.OperationResults),
						"statusCode", result.StatusCode,
						"resourceBody", string(result.ResourceBody),
						"transaction", t.ToJSONStructRendering(),
					)
					return nil, cosmosstorageutils.NewTransactionStepError(step+1, len(response.OperationResults), int(result.StatusCode), t.stepsDetails[step])
				}
			}
		}

		// The two slices SHOULD be of equal length.
		safeStop := min(len(itemIDs), len(response.OperationResults))
		for i := 0; i < safeStop; i++ {
			if len(response.OperationResults[i].ResourceBody) > 0 {
				result.items[itemIDs[i]] = response.OperationResults[i].ResourceBody
			}
		}
	}

	for _, callback := range t.onSuccess {
		callback(result)
	}

	return result, nil
}

func (t *cosmosDBTransaction) String() string {
	ret, err := json.Marshal(t.ToJSONStructRendering())
	if err != nil {
		return "failed to marshal transaction details: " + err.Error()
	}
	return string(ret)
}

// ToJSONStructRendering is useful for logging.
func (t *cosmosDBTransaction) ToJSONStructRendering() cosmosstorageutils.CosmosDBTransactionDetails {
	return cosmosstorageutils.CosmosDBTransactionDetails{
		PartitionKey: t.pk,
		Steps:        t.stepsDetails,
	}
}

var _ cosmosstorageutils.DBTransactionResult = &cosmosDBTransactionResult{}

type cosmosDBTransactionResult struct {
	items map[string]json.RawMessage
}

func newCosmosDBTransactionResult() *cosmosDBTransactionResult {
	return &cosmosDBTransactionResult{make(map[string]json.RawMessage)}
}

func getCastResult[InternalAPIType, CosmosAPIType any](r *cosmosDBTransactionResult, cosmosUID string) (*InternalAPIType, error) {
	data, ok := r.items[cosmosUID]
	if !ok {
		return nil, ErrItemNotFound
	}

	var cosmosObj CosmosAPIType
	if err := json.Unmarshal(data, &cosmosObj); err != nil {
		return nil, fmt.Errorf("failed to unmarshal Cosmos DB item '%s': %w", cosmosUID, err)
	}

	return cosmosstorageutils.CosmosToInternal[InternalAPIType, CosmosAPIType](&cosmosObj)
}

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

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

	switch strings.ToLower(typedDoc.ResourceType) {
	case strings.ToLower(coreapi.ClusterResourceType.String()):
		return getCastResult[coreapi.HCPOpenShiftCluster, cosmosstorageutils.GenericDocument[coreapi.HCPOpenShiftCluster]](r, cosmosUID)
	case strings.ToLower(coreapi.NodePoolResourceType.String()):
		return getCastResult[coreapi.HCPOpenShiftClusterNodePool, cosmosstorageutils.GenericDocument[coreapi.HCPOpenShiftClusterNodePool]](r, cosmosUID)
	case strings.ToLower(coreapi.ExternalAuthResourceType.String()):
		return getCastResult[coreapi.HCPOpenShiftClusterExternalAuth, cosmosstorageutils.GenericDocument[coreapi.HCPOpenShiftClusterExternalAuth]](r, cosmosUID)
	default:
		return nil, fmt.Errorf("unknown resource type '%s'", typedDoc.ResourceType)
	}
}
