package handshake

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/pingcap/tiproxy/pkg/proxy/backend"
	pnet "github.com/pingcap/tiproxy/pkg/proxy/net"
	"github.com/pingcap/tiproxy/pkg/proxy/proxyprotocol"
	"github.com/pkg/errors"
	"github.com/tidbcloud/aws-shared-provider/pkg/gateway/config"
	"github.com/tidbcloud/aws-shared-provider/pkg/gateway/meta"
	"github.com/tidbcloud/aws-shared-provider/pkg/gateway/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/gateway/probe"
	"github.com/tidbcloud/aws-shared-provider/pkg/infra"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidb"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	serverlessinfra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	ctrl "sigs.k8s.io/controller-runtime"
)

const (
	trafficRecordInterval = time.Minute

	DefaultAttributesValue = "unknown"
	ErrUserPrefixMismatch  = "ERROR 20003"
	AttrsConnKey           = "TiDB-Gateway-ConnID"

	connectWrongTiDBBackoff = 3 * time.Second
	connectWrongTiDBMaxWait = 30 * time.Second

	checkTidbConnTimeout = 7 * time.Second

	// AWS Proxy Protocol v2 TLV types
	PP2_TYPE_AWS            = 0xEA
	PP2_SUBTYPE_AWS_VPCE_ID = 0x01
	// Alibaba Cloud Proxy Protocol v2 TLV type
	PP2_TYPE_ALIBABA             = 0xE1
	PP2_SUBTYPE_ALIBABA_VPC_ID   = 0x01
	PP2_SUBTYPE_ALIBABA_EP_ID    = 0x02
	PP2_SUBTYPE_ALIBABA_EPSVC_ID = 0x03
)

func (h *HsHandler) getBackendErrType(cctx *ConnContext, backendAddr string) (string, error) {
	backendIP, _, err := net.SplitHostPort(backendAddr)
	if err != nil {
		return "", err
	}

	switch cctx.clusterType {
	case infra.TypeCluster:
		infraCluster, err := h.regionalInfra.GetClusterWithToken(h.ctx, cctx.token)
		if err != nil {
			if !errors.Is(err, infra.ErrNotFound) {
				return "", err
			} else {
				return config.CONNECTION_STATUS_DESIGN_ERROR, nil
			}
		}
		if infraCluster.DeletionTimestamp == nil && !infra.IsClusterRestoring(infraCluster) {
			if h.checkTidbConnNormalClosed(cctx, backendIP) {
				return config.CONNECTION_STATUS_DESIGN_ERROR, nil
			} else {
				return config.CONNECTION_STATUS_INNER_ERROR, nil
			}
		} else {
			return config.CONNECTION_STATUS_DESIGN_ERROR, nil
		}
	case infra.TypeBranch:
		branchCluster, err := h.infraApi.GetBranchWithToken(h.ctx, cctx.token)
		if err != nil {
			if !errors.Is(err, infra.ErrNotFound) {
				return "", err
			} else {
				return config.CONNECTION_STATUS_DESIGN_ERROR, nil
			}
		}
		if branchCluster.DeletionTimestamp == nil && !infra.IsBranchRestoring(branchCluster) {
			if h.checkTidbConnNormalClosed(cctx, backendIP) {
				return config.CONNECTION_STATUS_DESIGN_ERROR, nil
			} else {
				return config.CONNECTION_STATUS_INNER_ERROR, nil
			}
		} else {
			return config.CONNECTION_STATUS_DESIGN_ERROR, nil
		}
	default:
		return "", errors.Errorf("%s unknown cluster type %s", cctx.clusterID, cctx.clusterType)
	}
}

func (h *HsHandler) recordFlowBytes(cctx *ConnContext, in, out uint64) {
	if in > cctx.clientInBytes {
		metrics.ConnectionFlow.WithLabelValues(cctx.clusterID, cctx.keyspaceID, "in").Add(float64(in - cctx.clientInBytes)) //nolint:lll // ignore
		cctx.clientInBytes = in
	}
	if out > cctx.clientOutBytes {
		metrics.ConnectionFlow.WithLabelValues(cctx.clusterID, cctx.keyspaceID, "out").Add(float64(out - cctx.clientOutBytes)) //nolint:lll // ignore
		cctx.clientOutBytes = out
	}
}

func (h *HsHandler) setTLSState(res *pnet.HandshakeResp, ctx backend.ConnContext) {
	const secureConnKey = "TiDB-Gateway-Secure-Connection"
	if s := ctx.Value(backend.ConnContextKeyTLSState); s != nil {
		state := s.(tls.ConnectionState)
		res.Attrs[secureConnKey] = fmt.Sprintf(`{"Version":%d,"CipherSuite":%d}`, state.Version, state.CipherSuite)
	} else {
		// reset attr in case of evil client.
		delete(res.Attrs, secureConnKey)
	}
}

//nolint:mnd // ignore
func (h *HsHandler) parseToken(cctx *ConnContext, res *pnet.HandshakeResp) (string, error) {
	logger := ctrl.Log.WithName("parseToken").WithValues("connID", cctx.connID).WithValues("user_name", res.User)
	// user name format 1 (default): {token}.{uname}, demo: 38FwWRFUHqwi73c.root
	splits := strings.SplitN(res.User, ".", 2)
	if len(splits) != 2 {
		return "", utils.NewExternalError(utils.ERR_USER_MISSING_TOKEN)
	}
	if len(splits[0]) == utils.GatewayTokenLen {
		/* NLB/Privatelink prohibited cloud_admin user */
		if splits[1] == "cloud_admin" &&
			(h.conf.GatewayMode == config.PublicMode || h.conf.GatewayMode == config.PrivatelinkMode) &&
			!h.conf.E2EEnv {
			return "", utils.NewExternalError(fmt.Sprintf(utils.ERR_FORBIDDEN_USER, res.User))
		}
		// format 1
		return splits[0], nil
	}
	if h.conf.GatewayMode != config.InnerMode {
		return "", utils.NewExternalError(utils.ERR_CLUSTER_NOT_FOUND)
	}

	// try to parse format 2, only for config.InnerMode
	// user name format 2 (compatible with inner service): {namespace}.{name}.{uname},
	// demo: nightly-ms-43549.62791.root
	splits = strings.SplitN(res.User, ".", 3)
	if len(splits) != 3 {
		return "", utils.NewExternalError(utils.ERR_USER_MISSING_TOKEN)
	}
	logger.Info("(inner gateway) try to translate token", "input", res.User)

	isBranch := false
	var infraBranch *cloudv1.Branch
	infraCluster, err := h.regionalInfra.GetCluster(h.ctx, splits[0], splits[1])
	if err != nil {
		infraBranch, err = h.infraApi.GetBranch(h.ctx, splits[0], splits[1])
		// try to get branch from infra
		if err != nil {
			logger.Error(err, "get infra cluster by namespaced name failed.")
			return "", errors.New("can not find infra cluster by identifier")
		}
		isBranch = true
	}

	if !isBranch {
		res.User = infraCluster.Status.KeyspaceToken + "." + splits[2]
		return infraCluster.Status.KeyspaceToken, nil
	}

	clusterStatus, err := infraBranch.AwsFreeTierClusterStatusFromBranch()
	if err != nil {
		logger.Error(err, "get infra cluster status failed.")
		/* error message used for SLA statistics */
		return "", errors.New("can not get token from infra")
	}
	token := clusterStatus.KeyspaceToken
	res.User = token + "." + splits[2]
	return token, nil
}

func (h *HsHandler) checkConnLimit(cctx *ConnContext) error {
	if strings.Contains(cctx.tidbPool, "super-vip") {
		return nil
	}

	cfg := config.GlobalDynamicConfig()
	if cfg == nil {
		return nil
	}
	if !cfg.ConnLimit.Enable {
		return nil
	}

	/* limit maximum number of concurrent connections */
	var maxConnections int64
	if limit, ok := cfg.ConnLimit.ClusterSpecific[cctx.clusterID]; ok {
		maxConnections = limit
	} else if strings.Contains(cctx.tidbPool, "vip") {
		maxConnections = cfg.ConnLimit.PaidStarter
	} else {
		maxConnections = cfg.ConnLimit.FreeStarter
	}

	// Use tiproxy router's connection count (updated by tiproxy on connect/disconnect).
	// When no router exists yet (first connection for this token), count is 0.
	cc := int64(h.routerBundle.GetConnCountByToken(cctx.token))
	if cc > maxConnections {
		ctrl.Log.WithName("checkConnLimit").WithValues("connID", cctx.connID).
			WithValues("token", cctx.token).WithValues("cluster_id", cctx.clusterID).
			WithValues("cc", cc).WithValues("maxConns", maxConnections).Info("connection limit exceeded")
		// Temporarily disable connection limitation
		// return utils.NewExternalError(utils.ERR_CONNECTION_LIMIT_EXCEEDED)
	}
	return nil
}

func (h *HsHandler) checkInfraClusterStatus(cctx *ConnContext) error {
	meta, err := meta.GetClusterOrBranchMeta(h.ctx, h.regionalInfra, h.infraApi, h.conf, cctx.token)
	if err != nil {
		return err
	}
	cctx.clusterType = meta.ClusterType
	cctx.enableAuditLog = meta.EnableAuditLog
	cctx.auditLogRedacted = meta.AuditLogRedacted
	cctx.zoneRedundantEnabled = meta.ZoneRedundantEnabled
	cctx.keyspaceID = meta.KeyspaceID
	cctx.ruMode = meta.RuMode
	cctx.tidbPool = meta.TiDBPool
	cctx.clusterID = meta.ClusterID
	cctx.tenantID = meta.TenantID
	cctx.projectID = meta.ProjectID
	cctx.servicePlan = meta.ServicePlan

	if h.conf.GatewayMode == config.PublicMode {
		clientIP := strings.Split(cctx.clientAddr, ":")[0]
		if clientIP != "" && !utils.ValidAuthorizedNetworks(meta.AuthorizedNetworks, clientIP) {
			ctrl.Log.WithName("authorized_networks").
				WithValues("connID", cctx.connID).
				WithValues("token", cctx.token).
				WithValues("configured_authorized_networks", formatAuthorizedNetworksForLog(meta.AuthorizedNetworks)).
				Info(fmt.Sprintf(utils.ERR_CLUSTER_BLOCK_BY_TRAFFIC_FILTER, clientIP))
			return utils.NewExternalError(fmt.Sprintf(utils.ERR_CLUSTER_BLOCK_BY_TRAFFIC_FILTER, clientIP, cctx.clusterID))
		}
	}
	if h.conf.GatewayMode == config.PrivatelinkMode && !utils.ValidPrivatelinkAuthorizedNetworks(meta.PrivatelinkAuthorizedNetworks, cctx.vpcEndpointID) {
		ctrl.Log.WithName("privatelink_authorized_endpoints").
			WithValues("connID", cctx.connID).
			WithValues("token", cctx.token).
			WithValues("vpce", cctx.vpcEndpointID).
			WithValues("configured_privatelink_authorized_networks", formatPrivatelinkAuthorizedNetworksForLog(meta.PrivatelinkAuthorizedNetworks)).
			Info(fmt.Sprintf(utils.ERR_PRIVATELINK_ENDPOINT_BLOCKED, cctx.vpcEndpointID))
		return utils.NewExternalError(fmt.Sprintf(utils.ERR_PRIVATELINK_ENDPOINT_BLOCKED, cctx.vpcEndpointID, cctx.clusterID))
	}

	return nil
}

func (h *HsHandler) checkBackend(cctx *ConnContext, addr string, src *backend.ErrorSource) {
	logger := ctrl.Log.WithName("checkBackend").WithValues("connID", cctx.connID).
		WithValues("token", cctx.token).WithValues("cluster_id", cctx.clusterID)
	logger.Info("try to check backend", "addr", addr, "errorSource", src)

	// addr can be empty if tiproxy does not find any instance.
	// if src isn't from backend, no need to check backend.
	// develop mode does not need to check backend.
	if addr == "" || (src != nil && src.GetSourceComp() != backend.CompBackend) || h.conf.DevelopMode {
		return
	}

	ip, _, err := net.SplitHostPort(addr)
	if err != nil {
		logger.Error(err, "SplitHostPort failed", "addr", addr)
		return
	}
	if tidbPool := h.routerBundle.GetBackendTidbPool(cctx.token, ip); tidbPool != "" {
		h.checkBackendServer.InsertTask(probe.CheckBackend{Token: cctx.token, Backend: ip, TidbPoolNs: tidbPool})
	}
}

func (h *HsHandler) checkTidbConnNormalClosed(cctx *ConnContext, backendIP string) bool {
	logger := ctrl.Log.WithName("checkTidbConnNormalClosed").WithValues("connID", cctx.connID).
		WithValues("token", cctx.token).WithValues("cluster_id", cctx.clusterID).
		WithValues("backendIP", backendIP)

	startTime := time.Now()
	ctx, cancel := context.WithTimeout(h.ctx, checkTidbConnTimeout)
	defer cancel()
	url := fmt.Sprintf("https://manager-server:8000/api/tidb/checkconn?token=%s&conn_id=%d&ip=%s",
		cctx.token, cctx.connID, backendIP)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		logger.Error(err, "failed to create request")
		return false
	}
	resp, err := h.httpClient.Do(req)
	if err != nil {
		logger.Error(err, "failed to send request", "duration", time.Since(startTime).Seconds())
		return false
	}
	defer resp.Body.Close() //nolint:errcheck // no need check error
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		logger.Error(err, "failed to read response body", "duration", time.Since(startTime).Seconds())
		return false
	}
	if resp.StatusCode != http.StatusOK {
		logger.Error(nil, "manager return unexpected",
			"status", resp.Status, "body", string(body), "duration", time.Since(startTime).Seconds())
		return false
	}

	logger.Info("CheckConn", "msg", string(body), "duration", time.Since(startTime).Seconds())
	return string(body) == tidb.ConnNormalClosed
}

func (h *HsHandler) parseCloudTlv(proxy *proxyprotocol.Proxy, typ, subType byte) string {
	if proxy == nil {
		return ""
	}
	for _, tlv := range proxy.TLV {
		if tlv.Typ == proxyprotocol.ProxyTlvType(typ) && len(tlv.Content) > 0 && tlv.Content[0] == subType {
			return string(tlv.Content[1:])
		}
	}
	return ""
}

func (h *HsHandler) parsePrivatelinkEndpointID(proxy *proxyprotocol.Proxy) string {
	if proxy == nil || h.conf == nil {
		return ""
	}
	switch h.conf.CloudProviderType {
	case types.ProviderTypeAWSFreeTier:
		return h.parseCloudTlv(proxy, PP2_TYPE_AWS, PP2_SUBTYPE_AWS_VPCE_ID)
	case types.ProviderTypeAliCloudFreeTier:
		return h.parseCloudTlv(proxy, PP2_TYPE_ALIBABA, PP2_SUBTYPE_ALIBABA_EP_ID)
	default:
		return ""
	}
}

func formatAuthorizedNetworksForLog(networks []serverlessinfra.AuthorizedNetwork) []string {
	formatted := make([]string, 0, len(networks))
	for _, network := range networks {
		formatted = append(formatted, fmt.Sprintf("%s-%s", network.StartIpAddress, network.EndIpAddress))
	}
	return formatted
}

func formatPrivatelinkAuthorizedNetworksForLog(networks []serverlessinfra.PrivatelinkAuthorizedNetwork) []string {
	formatted := make([]string, 0, len(networks))
	for _, network := range networks {
		formatted = append(formatted, network.VpcEndpointId)
	}
	return formatted
}

// shouldSkipTLSCheck checks if a user should be exempt from TLS requirement
// based on username prefix matching.
func shouldSkipTLSCheck(username string, skipPrefixes []string) bool {
	if username == "" {
		return false
	}
	for _, rawPrefix := range skipPrefixes {
		prefix := strings.TrimSpace(rawPrefix)
		if prefix == "" {
			continue
		}
		if strings.HasPrefix(username, prefix) {
			return true
		}
	}
	return false
}