// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may not
// use this file except in compliance with the License. A copy of the
// License is located at
//
// http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 endpoint

import (
	"regexp"
	"strings"
	"sync"

	"github.com/aws/amazon-ssm-agent/agent/appconfig"
	"github.com/aws/amazon-ssm-agent/agent/log"
)

// Map defining region prefixes and the default service domain for prefix, keys are processed in random order
var regionPrefixServiceDomain = map[string]string{
	"cn-":      "amazonaws.com.cn",
	"us-iso-":  "c2s.ic.gov",
	"us-isob-": "sc2s.sgov.gov",
	"us-isof-": "csp.hci.ic.gov",
	"eu-isoe-": "cloud.adc-e.uk",
	"eusc-de-": "amazonaws.eu",
}

// Map defining region prefixes and dual-stack service domains
var regionPrefixServiceDualStackDomain = map[string]string{
	"cn-":      "api.amazonwebservices.com.cn",
	"us-iso-":  "api.aws.ic.gov",
	"us-isob-": "api.aws.scloud",
	"us-isof-": "api.aws.hci.ic.gov",
	"eu-isoe-": "api.cloud-aws.adc-e.uk",
	"eusc-de-": "api.amazonwebservices.eu",
}

// default service domain if prefix does not exist in awsFallbackServiceDomain map
const (
	defaultServiceDomain          = "amazonaws.com"
	defaultServiceDualStackDomain = "api.aws"
	s3Service                     = "s3"
	s3DualStackEndpointPrefix     = "s3.dualstack."
	regionMaxLength               = 100
)

// The following regex only allows a-z upper/lower case characters, digits and dashes for region strings
var supportedRegionReg = regexp.MustCompile(`^[A-Za-z0-9-]+$`)

type IEndpointHelper interface {
	// GetServiceEndpoint returns the endpoint for a service in a specific region
	GetServiceEndpoint(service, region string) string
}

func GetServiceDomainByPrefix(region string) string {
	for regionPrefix, serviceDomain := range regionPrefixServiceDomain {
		if strings.HasPrefix(region, regionPrefix) {
			return serviceDomain
		}
	}

	return defaultServiceDomain
}

func GetServiceDualStackDomainByPrefix(service, region string) string {
	// S3 uses regular service domains in dualstack endpoints
	if service == s3Service {
		return GetServiceDomainByPrefix(region)
	}

	// All other services use dedicated domains in dualstack endpoints
	for regionPrefix, dualStackDomain := range regionPrefixServiceDualStackDomain {
		if strings.HasPrefix(region, regionPrefix) {
			return dualStackDomain
		}
	}

	return defaultServiceDualStackDomain
}

type endpointImpl struct {
	log    log.T
	config appconfig.SsmagentConfig

	cacheLock                  sync.RWMutex
	regionServiceEndpointCache map[string]map[string]string
}

func (e *endpointImpl) isRegionValid(region string) bool {
	return len(region) <= regionMaxLength && supportedRegionReg.Match([]byte(region))
}

func (e *endpointImpl) endpointCacheLookup(service, region string) string {
	e.cacheLock.RLock()
	defer e.cacheLock.RUnlock()
	if serviceEndpointCache, ok := e.regionServiceEndpointCache[region]; ok {
		if endpoint, ok := serviceEndpointCache[service]; ok {
			return endpoint
		}
	}
	return ""
}

func (e *endpointImpl) setEndpointCache(service, region, endpoint string) {
	e.cacheLock.Lock()
	defer e.cacheLock.Unlock()
	if _, ok := e.regionServiceEndpointCache[region]; !ok {
		e.regionServiceEndpointCache[region] = map[string]string{}
	}

	e.regionServiceEndpointCache[region][service] = endpoint
}

func (e *endpointImpl) GetServiceEndpoint(service, region string) string {
	e.log.Debugf("Determining endpoint for service %s in region %s", service, region)

	if region == "" {
		// If region is not defined, we are unable to determine endpoint for the service
		e.log.Errorf("Cannot get endpoint for service %s due to unspecified region.", service)
		return ""
	} else if service == "" {
		// If service is not defined, we are unable to determine endpoint
		e.log.Errorf("Cannot get endpoint for service in region %s due to unspecified service.", service)
		return ""
	} else if !e.isRegionValid(region) {
		e.log.Errorf("Unable to determine endpoint because region %s has invalid characters", region)
		return ""
	} else if endpoint := e.endpointCacheLookup(service, region); endpoint != "" {
		return endpoint
	}

	var endpoint string

	// Check if dual-stack endpoints are enabled
	if e.config.Agent.UseDualStackEndpoint {
		endpoint = e.buildDualStackEndpoint(service, region)
	} else {
		endpoint = e.buildStandardEndpoint(service, region)
	}

	e.setEndpointCache(service, region, endpoint)
	return endpoint
}

func (e *endpointImpl) buildDualStackEndpoint(service, region string) string {
	var serviceDomain string
	if e.config.Agent.ServiceDomain != "" {
		serviceDomain = e.config.Agent.ServiceDomain
	} else {
		serviceDomain = GetServiceDualStackDomainByPrefix(service, region)
	}

	// S3 has a special dual-stack endpoint format s3.dualstack.{region}.{regularServiceDomain}
	if service == s3Service {
		return s3DualStackEndpointPrefix + region + "." + serviceDomain
	}

	// All other services use the standard dual-stack format {service}.{region}.{dualStackServiceDomain}
	return service + "." + region + "." + serviceDomain
}

func (e *endpointImpl) buildStandardEndpoint(service, region string) string {
	var serviceDomain string
	if e.config.Agent.ServiceDomain != "" {
		serviceDomain = e.config.Agent.ServiceDomain
	} else {
		serviceDomain = GetServiceDomainByPrefix(region)
	}

	// Build the standard endpoint for the service in the region
	return service + "." + region + "." + serviceDomain
}

func NewEndpointHelper(log log.T, config appconfig.SsmagentConfig) *endpointImpl {
	return &endpointImpl{
		log,
		config,
		sync.RWMutex{},
		map[string]map[string]string{},
	}
}
