// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

package azcosmos

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	azlog "github.com/Azure/azure-sdk-for-go/sdk/azcore/log"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/streaming"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/tracing"
	"github.com/Azure/azure-sdk-for-go/sdk/internal/log"
	"github.com/Azure/azure-sdk-for-go/sdk/internal/uuid"
)

const (
	apiVersion = "2020-11-05"
)

// Client is used to interact with the Azure Cosmos DB database service.
type Client struct {
	endpoint    string
	internal    *azcore.Client
	gem         *globalEndpointManager
	endpointUrl *url.URL
	caches      *sharedCacheSet
	closeOnce   sync.Once
}

// getContainerCache returns the container properties cache for this client.
func (c *Client) getContainerCache() *containerPropertiesCache {
	if c.caches == nil {
		return nil
	}
	return c.caches.containerCache
}

// getPKRangeCache returns the partition key range cache for this client.
func (c *Client) getPKRangeCache() *partitionKeyRangeCache {
	if c.caches == nil {
		return nil
	}
	return c.caches.pkRangeCache
}

// Endpoint used to create the client.
func (c *Client) Endpoint() string {
	return c.endpoint
}

// Close releases the shared cache reference for this client. The underlying
// caches are removed from the global registry once all clients to the same
// account endpoint have been closed. After Close, the client should not be used.
// Close is idempotent; calling it multiple times is safe.
func (c *Client) Close() {
	c.closeOnce.Do(func() {
		if c.endpoint != "" {
			releaseCaches(c.endpoint)
		}
	})
}

// NewClientWithKey creates a new instance of Cosmos client with shared key authentication. It uses the default pipeline configuration.
// endpoint - The cosmos service endpoint to use.
// cred - The credential used to authenticate with the cosmos service.
// options - Optional Cosmos client options.  Pass nil to accept default values.
func NewClientWithKey(endpoint string, cred KeyCredential, o *ClientOptions) (*Client, error) {
	endpointUrl, err := url.Parse(endpoint)
	if err != nil {
		return nil, err
	}
	preferredRegions := []string{}
	enableCrossRegionRetries := true
	if o != nil {
		preferredRegions = o.PreferredRegions
	}
	o = withDefaultTransport(o)

	gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newSharedKeyCredPolicy(cred), o), preferredRegions, 0, enableCrossRegionRetries)
	if err != nil {
		return nil, err
	}

	internalClient, err := newClient(newSharedKeyCredPolicy(cred), gem, o)
	if err != nil {
		return nil, err
	}
	return &Client{endpoint: endpoint, endpointUrl: endpointUrl, internal: internalClient, gem: gem, caches: acquireCaches(endpoint)}, nil
}

// NewClient creates a new instance of Cosmos client with Azure AD access token authentication. It uses the default pipeline configuration.
// endpoint - The cosmos service endpoint to use.
// cred - The credential used to authenticate with the cosmos service.
// options - Optional Cosmos client options.  Pass nil to accept default values.
func NewClient(endpoint string, cred azcore.TokenCredential, o *ClientOptions) (*Client, error) {
	endpointUrl, err := url.Parse(endpoint)
	if err != nil {
		return nil, err
	}

	var scope []string

	if o != nil && o.Cloud.Services != nil {
		if svcCfg, ok := o.Cloud.Services[ServiceName]; ok && svcCfg.Audience != "" {
			audience := svcCfg.Audience
			scope = []string{audience + "/.default"}
			log.Write(azlog.EventRequest, fmt.Sprintf("Using custom scope for authentication: %s", scope[0]))
		}
	}

	if scope == nil {
		// Fallback to account-scope
		scope, err = createScopeFromEndpoint(endpointUrl)
		if err != nil {
			return nil, err
		}
		log.Write(azlog.EventRequest, fmt.Sprintf("Using account scope from endpoint for authentication: %s", scope[0]))
	}

	preferredRegions := []string{}
	enableCrossRegionRetries := true
	if o != nil {
		preferredRegions = o.PreferredRegions
	}
	o = withDefaultTransport(o)
	gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o), preferredRegions, 0, enableCrossRegionRetries)
	if err != nil {
		return nil, err
	}

	internalClient, err := newClient(newCosmosBearerTokenPolicy(cred, scope, nil), gem, o)
	if err != nil {
		return nil, err
	}
	return &Client{endpoint: endpoint, endpointUrl: endpointUrl, internal: internalClient, gem: gem, caches: acquireCaches(endpoint)}, nil
}

// NewClientFromConnectionString creates a new instance of Cosmos client from connection string. It uses the default pipeline configuration.
// connectionString - The cosmos service connection string.
// options - Optional Cosmos client options.  Pass nil to accept default values.
func NewClientFromConnectionString(connectionString string, o *ClientOptions) (*Client, error) {
	const (
		accountEndpoint = "AccountEndpoint"
		accountKey      = "AccountKey"
	)

	splits := strings.SplitN(connectionString, ";", 2)
	if len(splits) < 2 {
		return nil, errors.New("failed parsing connection string due to it not consist of two parts separated by ';'")
	}

	var endpoint string
	var cred KeyCredential
	for _, split := range splits {
		keyVal := strings.SplitN(split, "=", 2)
		if len(keyVal) < 2 {
			return nil, fmt.Errorf("failed parsing connection string due to unmatched key value separated by '='")
		}
		switch {
		case strings.EqualFold(accountEndpoint, keyVal[0]):
			endpoint = keyVal[1]
		case strings.EqualFold(accountKey, keyVal[0]):
			c, err := NewKeyCredential(strings.TrimSuffix(keyVal[1], ";"))
			if err != nil {
				return nil, err
			}
			cred = c
		}
	}

	return NewClientWithKey(endpoint, cred, o)
}

func newClient(authPolicy policy.Policy, gem *globalEndpointManager, options *ClientOptions) (*azcore.Client, error) {
	if options == nil {
		options = &ClientOptions{}
	}
	if options.Logging.AllowedHeaders == nil {
		options.Logging.AllowedHeaders = getAllowedHeaders()
	}
	return azcore.NewClient(moduleName, serviceLibVersion,
		azruntime.PipelineOptions{
			PerCall: []policy.Policy{
				&headerPolicies{
					enableContentResponseOnWrite: options.EnableContentResponseOnWrite,
					priorityLevel:                options.PriorityLevel,
					throughputBucket:             options.ThroughputBucket,
				},
				&globalEndpointManagerPolicy{gem: gem},
			},
			PerRetry: []policy.Policy{
				// clientRetryPolicy is listed before authPolicy so that
				// its internal retry loop (continue inside Do) naturally
				// re-invokes authPolicy via req.Next() on each attempt.
				// This refreshes the Authorization header from the
				// credential cache on every internal retry — important
				// for short-lived AAD/MSI tokens that may expire across
				// the cumulative same-region + cross-region retry
				// window. Cosmos master/resource keys are unaffected
				// because they are static.
				&clientRetryPolicy{gem: gem},
				authPolicy,
			},
			Tracing: azruntime.TracingOptions{
				Namespace: "Microsoft.DocumentDB",
			},
		},
		&options.ClientOptions)
}

func newInternalPipeline(authPolicy policy.Policy, options *ClientOptions) azruntime.Pipeline {
	if options == nil {
		options = &ClientOptions{}
	}
	if options.Logging.AllowedHeaders == nil {
		options.Logging.AllowedHeaders = getAllowedHeaders()
	}
	return azruntime.NewPipeline(moduleName, serviceLibVersion,
		azruntime.PipelineOptions{
			PerRetry: []policy.Policy{
				authPolicy,
			},
		},
		&options.ClientOptions)
}

// withDefaultTransport returns a *ClientOptions whose Transport is set to the
// Cosmos default HTTP client when the caller has not supplied one. Callers
// of NewClient*/NewClientFromConnectionString invoke this exactly once and
// reuse the result for both the user-facing and the global-endpoint-manager
// pipelines so options are normalized in a single place. The returned value
// is always non-nil; only the top-level Transport field is set on the clone,
// so the caller's *ClientOptions struct is never mutated, but slice fields
// such as PreferredRegions or PerCallPolicies still share backing arrays
// with the caller's struct and should not be mutated in place.
func withDefaultTransport(options *ClientOptions) *ClientOptions {
	if options == nil {
		return &ClientOptions{
			ClientOptions: azcore.ClientOptions{Transport: defaultCosmosHTTPClient},
		}
	}
	if options.Transport != nil {
		return options
	}
	clone := *options
	clone.Transport = defaultCosmosHTTPClient
	return &clone
}

func createScopeFromEndpoint(endpoint *url.URL) ([]string, error) {
	return []string{fmt.Sprintf("%s://%s/.default", endpoint.Scheme, endpoint.Hostname())}, nil
}

// NewDatabase returns a struct that represents a database and allows database level operations.
// id - The id of the database.
func (c *Client) NewDatabase(id string) (*DatabaseClient, error) {
	if id == "" {
		return nil, errors.New("id is required")
	}

	return newDatabase(id, c)
}

// NewContainer returns a struct that represents a container and allows container level operations.
// databaseId - The id of the database.
// containerId - The id of the container.
func (c *Client) NewContainer(databaseId string, containerId string) (*ContainerClient, error) {
	if databaseId == "" {
		return nil, errors.New("databaseId is required")
	}

	if containerId == "" {
		return nil, errors.New("containerId is required")
	}

	db, err := newDatabase(databaseId, c)
	if err != nil {
		return nil, err
	}

	return db.NewContainer(containerId)
}

// CreateDatabase creates a new database.
// ctx - The context for the request.
// databaseProperties - The definition of the database
// o - Options for the create database operation.
func (c *Client) CreateDatabase(
	ctx context.Context,
	databaseProperties DatabaseProperties,
	o *CreateDatabaseOptions) (DatabaseResponse, error) {
	var err error
	spanName, err := getSpanNameForDatabases(c.accountEndpointUrl(), operationTypeCreate, resourceTypeDatabase, databaseProperties.ID)
	if err != nil {
		return DatabaseResponse{}, err
	}
	ctx, endSpan := startSpan(ctx, spanName.name, c.internal.Tracer(), &spanName.options)
	defer func() { endSpan(err) }()

	if o == nil {
		o = &CreateDatabaseOptions{}
	}
	returnResponse := true
	h := &headerOptionsOverride{
		enableContentResponseOnWrite: &returnResponse,
	}

	operationContext := pipelineRequestOptions{
		resourceType:          resourceTypeDatabase,
		resourceAddress:       "",
		isWriteOperation:      true,
		headerOptionsOverride: h,
	}

	path, err := generatePathForNameBased(resourceTypeDatabase, "", true)
	if err != nil {
		return DatabaseResponse{}, err
	}

	azResponse, err := c.sendPostRequest(
		path,
		ctx,
		databaseProperties,
		operationContext,
		nil,
		o.ThroughputProperties.addHeadersToRequest)
	if err != nil {
		return DatabaseResponse{}, err
	}

	response, err := newDatabaseResponse(azResponse)
	return response, err
}

// NewQueryDatabasesPager executes query for databases.
// query - The SQL query to execute.
// o - Options for the operation.
func (c *Client) NewQueryDatabasesPager(query string, o *QueryDatabasesOptions) *azruntime.Pager[QueryDatabasesResponse] {
	queryOptions := &QueryDatabasesOptions{}
	if o != nil {
		originalOptions := *o
		queryOptions = &originalOptions
	}

	operationContext := pipelineRequestOptions{
		resourceType:    resourceTypeDatabase,
		resourceAddress: "",
	}

	path, _ := generatePathForNameBased(resourceTypeDatabase, operationContext.resourceAddress, true)

	return azruntime.NewPager(azruntime.PagingHandler[QueryDatabasesResponse]{
		More: func(page QueryDatabasesResponse) bool {
			return page.ContinuationToken != nil
		},
		Fetcher: func(ctx context.Context, page *QueryDatabasesResponse) (QueryDatabasesResponse, error) {
			var err error
			spanName, err := getSpanNameForClient(c.accountEndpointUrl(), operationTypeQuery, resourceTypeDatabase, c.accountEndpointUrl().Hostname())
			if err != nil {
				return QueryDatabasesResponse{}, err
			}
			ctx, endSpan := startSpan(ctx, spanName.name, c.internal.Tracer(), &spanName.options)
			defer func() { endSpan(err) }()
			if page != nil {
				if page.ContinuationToken != nil {
					// Use the previous page continuation if available
					queryOptions.ContinuationToken = page.ContinuationToken
				}
			}

			azResponse, err := c.sendQueryRequest(
				path,
				ctx,
				query,
				queryOptions.QueryParameters,
				operationContext,
				queryOptions,
				nil)

			if err != nil {
				return QueryDatabasesResponse{}, err
			}

			return newDatabasesQueryResponse(azResponse)
		},
	})
}

func (c *Client) sendPostRequest(
	path string,
	ctx context.Context,
	content interface{},
	operationContext pipelineRequestOptions,
	requestOptions cosmosRequestOptions,
	requestEnricher func(*policy.Request)) (*http.Response, error) {
	req, err := c.createRequest(path, ctx, http.MethodPost, operationContext, requestOptions, requestEnricher)
	if err != nil {
		return nil, err
	}

	err = c.attachContent(content, req)
	if err != nil {
		return nil, err
	}

	return c.executeAndEnsureSuccessResponse(ctx, req)
}

func (c *Client) sendQueryRequest(
	path string,
	ctx context.Context,
	query string,
	parameters []QueryParameter,
	operationContext pipelineRequestOptions,
	requestOptions cosmosRequestOptions,
	requestEnricher func(*policy.Request)) (*http.Response, error) {
	req, err := c.createRequest(path, ctx, http.MethodPost, operationContext, requestOptions, requestEnricher)
	if err != nil {
		return nil, err
	}

	err = azruntime.MarshalAsJSON(req, queryBody{
		Query:      query,
		Parameters: parameters,
	})

	if err != nil {
		return nil, err
	}

	req.Raw().Header.Add(cosmosHeaderQuery, "True")
	// Override content type for query
	req.Raw().Header.Set(headerContentType, cosmosHeaderValuesQuery)

	return c.executeAndEnsureSuccessResponse(ctx, req)
}

func (c *Client) sendPutRequest(
	path string,
	ctx context.Context,
	content interface{},
	operationContext pipelineRequestOptions,
	requestOptions cosmosRequestOptions,
	requestEnricher func(*policy.Request)) (*http.Response, error) {
	req, err := c.createRequest(path, ctx, http.MethodPut, operationContext, requestOptions, requestEnricher)
	if err != nil {
		return nil, err
	}

	err = c.attachContent(content, req)
	if err != nil {
		return nil, err
	}

	return c.executeAndEnsureSuccessResponse(ctx, req)
}

func (c *Client) sendGetRequest(
	path string,
	ctx context.Context,
	operationContext pipelineRequestOptions,
	requestOptions cosmosRequestOptions,
	requestEnricher func(*policy.Request)) (*http.Response, error) {
	req, err := c.createRequest(path, ctx, http.MethodGet, operationContext, requestOptions, requestEnricher)
	if err != nil {
		return nil, err
	}

	return c.executeAndEnsureSuccessResponse(ctx, req)
}

func (c *Client) sendDeleteRequest(
	path string,
	ctx context.Context,
	operationContext pipelineRequestOptions,
	requestOptions cosmosRequestOptions,
	requestEnricher func(*policy.Request)) (*http.Response, error) {
	req, err := c.createRequest(path, ctx, http.MethodDelete, operationContext, requestOptions, requestEnricher)
	if err != nil {
		return nil, err
	}

	return c.executeAndEnsureSuccessResponse(ctx, req)
}

func (c *Client) sendBatchRequest(
	ctx context.Context,
	path string,
	batch []batchOperation,
	operationContext pipelineRequestOptions,
	requestOptions cosmosRequestOptions,
	requestEnricher func(*policy.Request)) (*http.Response, error) {
	req, err := c.createRequest(path, ctx, http.MethodPost, operationContext, requestOptions, requestEnricher)
	if err != nil {
		return nil, err
	}

	err = c.attachContent(batch, req)
	if err != nil {
		return nil, err
	}

	return c.executeAndEnsureSuccessResponse(ctx, req)
}

func (c *Client) sendPatchRequest(
	path string,
	ctx context.Context,
	content interface{},
	operationContext pipelineRequestOptions,
	requestOptions cosmosRequestOptions,
	requestEnricher func(*policy.Request)) (*http.Response, error) {
	req, err := c.createRequest(path, ctx, http.MethodPatch, operationContext, requestOptions, requestEnricher)
	if err != nil {
		return nil, err
	}

	err = c.attachContent(content, req)
	if err != nil {
		return nil, err
	}

	return c.executeAndEnsureSuccessResponse(ctx, req)
}

func (c *Client) createRequest(
	path string,
	ctx context.Context,
	method string,
	operationContext pipelineRequestOptions,
	requestOptions cosmosRequestOptions,
	requestEnricher func(*policy.Request)) (*policy.Request, error) {

	// todo: endpoint will be set originally by globalendpointmanager
	finalURL := c.endpoint

	if path != "" {
		finalURL = azruntime.JoinPaths(c.endpoint, path)
	}

	req, err := azruntime.NewRequest(ctx, method, finalURL)
	if err != nil {
		return nil, err
	}

	if requestOptions != nil {
		headers := requestOptions.toHeaders()
		if headers != nil {
			for k, v := range *headers {
				req.Raw().Header.Set(k, v)
			}
		}
	}

	addDefaultHeaders(req)

	req.SetOperationValue(operationContext)

	if requestEnricher != nil {
		requestEnricher(req)
	}

	req = attachRequestDiagnostics(req, operationContext)

	return req, nil
}

func (c *Client) attachContent(content interface{}, req *policy.Request) error {
	var err error
	switch v := content.(type) {
	case []byte:
		// If its a raw byte array, we can just set the body
		err = req.SetBody(streaming.NopCloser(bytes.NewReader(v)), "application/json")
	default:
		// Otherwise, we need to marshal it
		err = azruntime.MarshalAsJSON(req, content)
	}

	if err != nil {
		return err
	}

	return nil
}

func (c *Client) executeAndEnsureSuccessResponse(ctx context.Context, request *policy.Request) (*http.Response, error) {
	log.Write(azlog.EventResponse, fmt.Sprintf("\n===== Client preferred regions:\n%v\n=====\n", c.gem.preferredLocations))
	state := requestDiagnosticsStateFromContext(request.Raw().Context())
	finalizeRequestTrace := func() {
		if state != nil && state.requestTrace != nil {
			state.requestTrace.End()
		}
	}

	response, err := c.internal.Pipeline().Do(request)
	if err != nil {
		var responseErr *azcore.ResponseError
		if errors.As(err, &responseErr) && responseErr.RawResponse != nil {
			addPointOperationStatisticsFromResponse(responseErr.RawResponse, responseErr.Error(), traceDatumKeyPointOperationStatistics)
			finalizeRequestTrace()
			return nil, err
		}

		finalizeRequestTrace()
		return nil, wrapRequestError(err, diagnosticsFromContext(request.Raw().Context()))
	}

	c.addResponseValuesToSpan(ctx, response)

	successResponse := (response.StatusCode >= 200 && response.StatusCode < 300) || response.StatusCode == 304
	if successResponse {
		addPointOperationStatisticsFromResponse(response, "", traceDatumKeyPointOperationStatistics)
		finalizeRequestTrace()
		return response, nil
	}

	addPointOperationStatisticsFromResponse(response, response.Status, traceDatumKeyPointOperationStatistics)
	finalizeRequestTrace()
	return nil, azruntime.NewResponseErrorWithErrorCode(response, response.Status)
}

func (c *Client) accountEndpointUrl() *url.URL {
	return c.endpointUrl
}

func (c *Client) addResponseValuesToSpan(ctx context.Context, resp *http.Response) {
	span := c.internal.Tracer().SpanFromContext(ctx)
	span.SetAttributes(
		tracing.Attribute{Key: "db.cosmosdb.request_charge", Value: readRequestCharge(resp)},
		tracing.Attribute{Key: "db.cosmosdb.status_code", Value: resp.StatusCode},
	)
}

type pipelineRequestOptions struct {
	headerOptionsOverride *headerOptionsOverride
	resourceType          resourceType
	resourceAddress       string
	isRidBased            bool
	isWriteOperation      bool
}

func addDefaultHeaders(req *policy.Request) {
	req.Raw().Header.Set(headerXmsDate, time.Now().UTC().Format(http.TimeFormat))
	req.Raw().Header.Set(headerXmsVersion, apiVersion)
	req.Raw().Header.Set(cosmosHeaderSDKSupportedCapabilities, supportedCapabilitiesHeaderValue)
	if activityID, err := uuid.New(); err == nil {
		req.Raw().Header.Set(cosmosHeaderActivityId, activityID.String())
	}
}

func getAllowedHeaders() []string {
	headers := make([]string, 0, len(allKnownHttpHeaders))
	for _, h := range allKnownHttpHeaders {
		if h == headerAuthorization || h == "Proxy-Authorization" || h == "Transfer-Encoding" {
			continue
		}
		headers = append(headers, h)
	}
	return headers
}
