package handshake

import (
	"bytes"
	"context"
	"crypto/sha1"
	"fmt"
	"strings"
	"time"

	gomysql "github.com/go-mysql-org/go-mysql/mysql"
	tiproxyrouter "github.com/pingcap/tiproxy/pkg/manager/router"
	"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/client"
	"github.com/tidbcloud/aws-shared-provider/pkg/gateway/config"
	"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/gateway/router"
	"github.com/tidbcloud/aws-shared-provider/pkg/infra"
	"github.com/tidbcloud/aws-shared-provider/pkg/mysql"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"go.uber.org/zap"
	ctrl "sigs.k8s.io/controller-runtime"
)

// HsHandler implements tiproxy's backend.HandshakeHandler interface.
var _ backend.HandshakeHandler = &HsHandler{}

type HsHandler struct {
	ctx                context.Context
	conf               *config.Config
	regionalInfra      infra.RegionalInfra
	infraApi           infra.InfraAPI
	routerBundle       *router.RouterBundle
	zLogger            *zap.Logger
	httpClient         *client.ManagerServerClient
	probeTiDB          *probe.TiDBProber
	checkBackendServer *probe.CheckBackendServer
}

func NewHandshakeHandler(
	ctx context.Context,
	conf *config.Config,
	regionalInfra infra.RegionalInfra,
	infraApi infra.InfraAPI,
	routerBundle *router.RouterBundle,
	zLogger *zap.Logger,
	httpClient *client.ManagerServerClient,
	probeTiDB *probe.TiDBProber,
	checkBackendServer *probe.CheckBackendServer,
) backend.HandshakeHandler {
	return &HsHandler{
		ctx:                ctx,
		conf:               conf,
		regionalInfra:      regionalInfra,
		infraApi:           infraApi,
		routerBundle:       routerBundle,
		zLogger:            zLogger,
		httpClient:         httpClient,
		probeTiDB:          probeTiDB,
		checkBackendServer: checkBackendServer,
	}
}

func (h *HsHandler) GetCapability() pnet.Capability {
	return backend.SupportedServerCapabilities
}

func (h *HsHandler) GetServerVersion() string {
	return h.conf.TiDBServerVersion
}

func emptyPassowrdHash(seed []byte) ([]byte, error) {
	//https://dev.mysql.com/doc/dev/mysql-server/8.4.6/page_protocol_connection_phase_authentication_methods_native_password_authentication.html
	stage1 := sha1.Sum(nil)
	stage2 := sha1.Sum(stage1[:])
	h := sha1.New()
	_, err := h.Write(seed)
	if err != nil {
		return nil, err
	}
	_, err = h.Write(stage2[:])
	if err != nil {
		return nil, err
	}
	finalStage := h.Sum(nil)
	emptyPasswordHash := make([]byte, 20)
	for i := range emptyPasswordHash {
		emptyPasswordHash[i] = stage1[i] ^ finalStage[i]
	}
	return emptyPasswordHash, nil
}

func (h *HsHandler) HandleHandshakeResp(ctx backend.ConnContext, res *pnet.HandshakeResp) error {
	cctx := &ConnContext{
		user:              res.User,
		connID:            ctx.Value(backend.ConnContextKeyConnID).(uint64),
		start:             time.Now(),
		phase:             utils.GATEWAY_SESSION_PHASE_CONNECT_ACCEPT,
		status:            utils.GATEWAY_SESSION_STATUS_FAILED,
		lastTrafficRecord: time.Now(),

		clientName:    DefaultAttributesValue,
		clientVersion: DefaultAttributesValue,
		programName:   DefaultAttributesValue,
		clientAddr:    ctx.ClientAddr(), // ip:port "35.83.90.11:42480"
	}
	ctx.SetValue(connContextKey, cctx)

	if h.conf.DevelopMode {
		cctx.token = "DEV"
		return nil
	}

	metrics.ConnectionCounter.WithLabelValues(cctx.clusterID, utils.GATEWAY_SESSION_PHASE_CONNECT_ACCEPT).Inc()

	// TODO: set keepalive and nodelay options
	logger := ctrl.Log.WithName("HandleHandshakeResp").WithValues("connID", cctx.connID)
	logger.V(1).Info("accepting new connection")

	if h.conf.GatewayMode == config.PrivatelinkMode {
		if proxyProtocolInfo, ok := ctx.Value(backend.ConnContextKeyProxyProtocol).(*proxyprotocol.Proxy); ok {
			cctx.vpcEndpointID = h.parsePrivatelinkEndpointID(proxyProtocolInfo)
			logger = logger.WithValues("vpce", cctx.vpcEndpointID)
		}
	}

	// Serverless must use TLS
	if uint32(res.Capability)&mysql.ClientSSL == 0 && h.conf.TLSOnly {
		if !shouldSkipTLSCheck(res.User, h.conf.TLSOnlySkipUserPrefixes) {
			logger.Info("Connections using insecure transport are prohibited on serverless tier.")
			return utils.NewExternalError(utils.ERR_TLS_ONLY)
		}
		logger.Info("User is exempt from TLS requirement based on prefix", "user", res.User)
	}

	cctx.phase = utils.GATEWAY_SESSION_PHASE_TOKEN_PARSER
	startTime := time.Now()
	token, err := h.parseToken(cctx, res)
	if err != nil {
		return err
	}
	cctx.token = token
	logger = logger.WithValues("token", cctx.token)

	if err = h.checkInfraClusterStatus(cctx); err != nil {
		logger.Error(errors.Cause(err), "check infra cluster status failed")
		return err
	}
	logger.WithValues("cluster_id", cctx.clusterID)

	if h.conf.OnlyForZR && !cctx.zoneRedundantEnabled {
		return utils.NewExternalError(fmt.Sprintf(utils.ERR_PRIVATELINK_ENDPOINT_ZR_DISABLE, cctx.clusterID))
	}

	if err = h.checkConnLimit(cctx); err != nil {
		logger.Error(err, "free cluster limitations")
		return err
	}

	metrics.ConnectionCounter.WithLabelValues(cctx.clusterID, cctx.phase).Inc()
	metrics.ConnectionHistogram.WithLabelValues(cctx.clusterID, cctx.phase).Observe(time.Since(startTime).Seconds()) //nolint:lll // ignore
	cctx.phase = utils.GATEWAY_SESSION_PHASE_CONNECT_BACKEND

	if res.Attrs == nil {
		res.Attrs = make(map[string]string)
	}
	res.Attrs[AttrsConnKey] = fmt.Sprint(cctx.connID)

	// collect connection attributes
	if val, ok := res.Attrs["_client_name"]; ok {
		cctx.clientName = val
	}
	if val, ok := res.Attrs["_client_version"]; ok {
		cctx.clientVersion = val
	}
	if val, ok := res.Attrs["program_name"]; ok {
		cctx.programName = val
	}
	cctx.attributes = fmt.Sprintf("%v", res.Attrs)
	metrics.ClientsDriverCounter.WithLabelValues(cctx.clientName, cctx.clientVersion, cctx.programName).Inc()

	// If the connection between gateway and client is secure, setup an attr to let TiDB know.
	h.setTLSState(res, ctx)

	// Do empty password detection for `mysql_native_password`
	if res.AuthPlugin == pnet.AuthNativePassword {
		seed, ok := ctx.Value(backend.ConnContextKeyAuthSalt).([20]byte)
		if ok && len(seed) == 20 {
			if len(res.AuthData) == 0 {
				cctx.anomalous = true
			} else {
				emptyPassword, err := emptyPassowrdHash(seed[:])
				if err == nil && bytes.Equal(emptyPassword, res.AuthData) {
					cctx.anomalous = true
				}
			}
		}
	}
	return nil
}

func (h *HsHandler) GetRouter(ctx backend.ConnContext, _ *pnet.HandshakeResp) (tiproxyrouter.Router, error) {
	val := ctx.Value(connContextKey)
	if val == nil {
		return nil, utils.NewInternalError(utils.ERR_CLUSTER_TEMPORARY_UNAVAILABLE, errors.New("connContext is nil"))
	}
	cctx := val.(*ConnContext)
	if !cctx.loggerInitialized {
		ctx.UpdateLogger(zap.String("token", cctx.token), zap.String("cluster_id", cctx.clusterID))
		cctx.loggerInitialized = true
	}

	logger := ctrl.Log.WithName("GetSharedRouter").WithValues("connID", cctx.connID).
		WithValues("token", cctx.token).WithValues("cluster_id", cctx.clusterID)
	rt := h.routerBundle.TryGetRouter(cctx.token)
	if rt != nil {
		logger.Info("get exist router success")
		return rt, nil
	}

	logger.Info("create new router")
	return h.routerBundle.CreateRouter(h.ctx, h.zLogger, cctx.token, cctx.clusterID, h.conf)
}

func (h *HsHandler) OnHandshake(ctx backend.ConnContext, to string, err error, src backend.ErrorSource) {
	val := ctx.Value(connContextKey)
	if val == nil {
		return
	}
	cctx := val.(*ConnContext)
	cctx.err = err
	logger := ctrl.Log.WithName("OnHandshake").
		WithValues("connID", cctx.connID).WithValues("errorSource", src).
		WithValues("token", cctx.token).WithValues("cluster_id", cctx.clusterID)
	proxyProtocolInfo, ok := ctx.Value(backend.ConnContextKeyProxyProtocol).(*proxyprotocol.Proxy)
	if ok {
		switch h.conf.CloudProviderType {
		case types.ProviderTypeAliCloudFreeTier:
			vpc := h.parseCloudTlv(proxyProtocolInfo, PP2_TYPE_ALIBABA, PP2_SUBTYPE_ALIBABA_VPC_ID)
			ep := h.parseCloudTlv(proxyProtocolInfo, PP2_TYPE_ALIBABA, PP2_SUBTYPE_ALIBABA_EP_ID)
			epsvc := h.parseCloudTlv(proxyProtocolInfo, PP2_TYPE_ALIBABA, PP2_SUBTYPE_ALIBABA_EPSVC_ID)
			logger = logger.WithValues("vpc", vpc).WithValues("ep", ep).WithValues("epsvc", epsvc)
		case types.ProviderTypeAWSFreeTier:
			vpce := cctx.vpcEndpointID
			if vpce == "" {
				vpce = h.parseCloudTlv(proxyProtocolInfo, PP2_TYPE_AWS, PP2_SUBTYPE_AWS_VPCE_ID)
			}
			logger = logger.WithValues("vpce", vpce)
		}
	}
	if err != nil {
		logger.Info("handshake failed", "err", err, "host", to)
		if to != "" {
			h.checkBackend(cctx, to, &src)
		}
		return
	}

	metrics.ConnectionCounter.WithLabelValues(cctx.clusterID, utils.GATEWAY_SESSION_PHASE_CONNECT_BACKEND).Inc()
	if cctx.status != utils.GATEWAY_SESSION_STATUS_SUCCESS {
		duration := time.Since(cctx.start).Seconds()
		logger.V(1).Info("handshake success", "host", to, "duration", duration)
		metrics.ConnectionHistogram.WithLabelValues(cctx.clusterID, utils.GATEWAY_SESSION_PHASE_CONNECT_BACKEND).Observe(duration) //nolint:lll // ignore
	} else {
		logger.V(1).Info("handshake success", "host", to)
	}
	cctx.phase = utils.GATEWAY_SESSION_PHASE_RELAY_DATA
	cctx.status = utils.GATEWAY_SESSION_STATUS_SUCCESS

	// update some ConnContext fields
	_ = h.checkInfraClusterStatus(cctx)

	// start to probe TiDB pod
	h.probeTiDB.Inc(cctx.token, cctx.clusterID, fmt.Sprint(cctx.connID))

	// If the connection is anomalous, log it
	if cctx.anomalous {
		logger.V(1).Info("A no-password login attempt successed", "host", to)
	}
}

func (h *HsHandler) OnTraffic(ctx backend.ConnContext) {
	val := ctx.Value(connContextKey)
	if val == nil {
		return
	}
	cctx := val.(*ConnContext)
	if time.Since(cctx.lastTrafficRecord) < trafficRecordInterval {
		return
	}
	h.recordFlowBytes(cctx, ctx.ClientInBytes(), ctx.ClientOutBytes())
	cctx.lastTrafficRecord = time.Now()
}

// isClientConnectionError returns true if err indicates the client connection failed (EOF, timeout,
// reset, TLS errors). Used to classify such closes as EXTER_WARN (client-side).
func isClientConnectionError(err error) bool {
	if err == nil {
		return false
	}
	s := err.Error()
	return strings.Contains(s, "failed to read the connection") ||
		strings.Contains(s, "i/o timeout") ||
		strings.Contains(s, "connection reset by peer") ||
		strings.Contains(s, "tls: user canceled") ||
		strings.Contains(s, "tls: received record with version") ||
		strings.Contains(s, "tls: bad record MAC")
}

func (h *HsHandler) OnConnClose(ctx backend.ConnContext, src backend.ErrorSource) error {
	val := ctx.Value(connContextKey)
	if val == nil {
		return nil
	}
	cctx := val.(*ConnContext)
	logger := ctrl.Log.WithName("OnConnClose").WithValues("error_source", src).WithValues("connID", cctx.connID).
		WithValues("token", cctx.token).WithValues("cluster_id", cctx.clusterID)

	metrics.ConnectionCounter.WithLabelValues(cctx.clusterID, utils.GATEWAY_SESSION_PHASE_CONNECT_CLOSE).Inc()
	metrics.SessionHistogram.WithLabelValues(cctx.clusterID).Observe(time.Since(cctx.start).Seconds()) //nolint:lll // ignore

	in, out := ctx.ClientInBytes(), ctx.ClientOutBytes()
	h.recordFlowBytes(cctx, in, out)

	errType := config.CONNECTION_STATUS_NORMAL
	switch src.GetSourceComp() {
	case backend.CompBackend:
		var err error
		if errType, err = h.getBackendErrType(cctx, ctx.ServerAddr()); err != nil {
			logger.Error(err, "failed to get backend error type")
			errType = config.CONNECTION_STATUS_INNER_ERROR
		} else if errType == config.CONNECTION_STATUS_INNER_ERROR {
			h.checkBackend(cctx, ctx.ServerAddr(), &src)
		}
	case backend.CompProxy:
		var externalErr *utils.ExternalError
		//nolint: gocritic // no need to use switch
		if src != backend.SrcProxyQuit && errors.As(cctx.err, &externalErr) {
			errType = config.CONNECTION_STATUS_EXTER_WARN
		} else if src == backend.SrcProxyNoBackend {
			// SrcProxyNoBackend may be caused by infra cluster deletion
			errType = config.CONNECTION_STATUS_INNER_ERROR
			if err := h.checkInfraClusterStatus(cctx); err != nil {
				errType = config.CONNECTION_STATUS_DESIGN_ERROR
			}
		} else {
			errType = config.CONNECTION_STATUS_INNER_ERROR
		}
	case backend.CompClient:
		if src == backend.SrcClientAuthFail && cctx.err != nil && strings.Contains(cctx.err.Error(), ErrUserPrefixMismatch) {
			// client error `ERROR 20003` is special, this is actually our problem
			errType = config.CONNECTION_STATUS_INNER_ERROR
		} else {
			errType = config.CONNECTION_STATUS_EXTER_WARN
		}
	default:
		// backend.CompNone here
		if cctx.err != nil {
			errType = config.CONNECTION_STATUS_INNER_WARN
		}
	}

	errMsg := ""
	if cctx.err != nil {
		errMsg = errors.Cause(cctx.err).Error()
	}

	clientErrMsg := ""
	clientErr, ok := ctx.Value(backend.ConnContextKeyClientError).(error)
	if ok && clientErr != nil {
		clientErrMsg = errors.Cause(clientErr).Error()
		// Prefer EXTER_WARN when context has a client connection error (same patterns as cctx.err),
		// but do not override backend-derived classification.
		if src.GetSourceComp() != backend.CompBackend &&
			errType != config.CONNECTION_STATUS_EXTER_WARN &&
			isClientConnectionError(clientErr) {
			errType = config.CONNECTION_STATUS_EXTER_WARN
		}
	}

	logger.Info("gw_access_log",
		"keyspace_id", cctx.keyspaceID,
		"phase", cctx.phase,
		"status", cctx.status,
		"duration", time.Since(cctx.start).Seconds(),
		"error_type", errType,
		"in_bytes", in,
		"out_bytes", out,
		"user", cctx.user,
		"tidb_pool", cctx.tidbPool,
		"ru_mode", cctx.ruMode,
		"backend_addr", ctx.ServerAddr(),
		"client_addr", ctx.ClientAddr(),
		"client_name", cctx.clientName,
		"client_version", cctx.clientVersion,
		"program_name", cctx.programName,
		"attributes", cctx.attributes,
		"connAddr", ctx.Value(backend.ConnContextKeyConnAddr),
		"error", errMsg,
		"client_error", clientErrMsg,
	)
	metrics.ConnectionErrorTypeCounter.WithLabelValues(cctx.tidbPool, cctx.ruMode, src.String(), errType).Inc()
	metrics.ConnectionSLACounter.WithLabelValues(cctx.clusterID, cctx.token, cctx.tidbPool, cctx.ruMode, src.String(), "connect", errType, cctx.servicePlan).Inc() //nolint:lll // ignore
	return nil
}

func (h *HsHandler) HandleHandshakeErr(ctx backend.ConnContext, err *gomysql.MyError) bool {
	val := ctx.Value(connContextKey)
	if val == nil {
		return false
	}
	cctx := val.(*ConnContext)

	logger := ctrl.Log.WithName("HandleHandshakeErr").WithValues("connID", cctx.connID).
		WithValues("token", cctx.token).WithValues("cluster_id", cctx.clusterID)

	if err != nil && strings.Contains(err.Error(), ErrUserPrefixMismatch) {
		logger.Info("user prefix mismatch, maybe connect wrong TiDB")
		h.checkBackend(cctx, ctx.ServerAddr(), nil)
		if time.Since(cctx.start) > connectWrongTiDBMaxWait {
			return false
		}
		time.Sleep(connectWrongTiDBBackoff)
		return true
	}

	return false
}