package httpgateway

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"net/http/httputil"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/pingcap/errors"
	_ "github.com/pingcap/tidb/pkg/types/parser_driver"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/log"
	"github.com/tidbcloud/pkgv2/logid"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

const (
	DatabaseHeader          = "TiDB-Database"
	SessionHeader           = "TiDB-Session"
	TxnIsolationLevelHeader = "TiDB-Isolation-Level"
	DebugTraceId            = "X-Debug-Trace-Id"
	StatefulActionHeader    = "TiDB-Stateful-Action"
	StatefulActionClose     = "close"
	StatefulActionOpen      = "open"
)

type HttpGateway struct {
	order         int
	componentName string
	poolManager   *PoolManager
	sessionStore  *StatefulSessionStore
}

func New(componentName string, logLevel string, secret string) (*HttpGateway, error) {
	splits := strings.Split(os.Getenv("HOSTNAME"), "-")
	order, err := strconv.Atoi(splits[len(splits)-1])
	if err != nil {
		return nil, fmt.Errorf("invalid pod host name: %s", os.Getenv("HOSTNAME"))
	}

	// set log level
	switch logLevel {
	case "DEBUG":
		log.SetLevel(zapcore.DebugLevel)
	case "INFO":
		log.SetLevel(zapcore.InfoLevel)
	case "WARN":
		log.SetLevel(zapcore.WarnLevel)
	default:
		log.SetLevel(zapcore.InfoLevel)
	}

	cipher, err := NewSessionCipher(secret)
	if err != nil {
		return nil, err
	}
	poolManager := NewPoolManager(cipher)
	sessionStore := NewStatefulSessionStore()

	return &HttpGateway{
		order:         order,
		componentName: componentName,
		poolManager:   poolManager,
		sessionStore:  sessionStore,
	}, nil
}

func (g *HttpGateway) ExecuteSQL(ctx *gin.Context) {
	// inject logid
	logId := ctx.GetHeader(DebugTraceId)
	if logId == "" {
		logId = logid.GenLogid()
		ctx.Header(DebugTraceId, logId)
	}
	ctxReq := logid.WithLogid(ctx.Request.Context(), logId)
	ctx.Request = ctx.Request.WithContext(ctxReq)

	startT := time.Now()
	defer func() {
		// collect metric
		status := ctx.Writer.Status()
		HTTPGatewayHistogram.WithLabelValues(strconv.Itoa(status)).Observe(time.Since(startT).Seconds())
	}()

	session := ctx.GetHeader(SessionHeader)
	logger := log.WithContextL(ctxReq).With(zap.String("session", session))
	logger.Info("receive request")

	// only stateful request with session may need to be forwarded
	if session != "" && !strings.Contains(session, StatelessSessionPrefix) {
		order, err := GetStatefulSessionOrder(session)
		if err != nil {
			logger.Error("failed to get session order", zap.Error(err))
			ctx.JSON(http.StatusBadRequest, NewErrorResp(http.StatusBadRequest, InvalidSessionErrorMessage))
			return
		}
		// forward to the right statefulSet pod
		if order != g.order {
			if err := g.forward(ctx, order); err != nil {
				logger.Error("failed to forward stateful session", zap.Error(err))
				ctx.JSON(http.StatusInternalServerError, NewErrorResp(http.StatusInternalServerError, UnknownInternalErrorMessage))
				return
			}
			return
		}
	}

	// check header
	action := strings.ToLower(ctx.GetHeader(StatefulActionHeader))
	if action == StatefulActionOpen && session != "" {
		ctx.JSON(http.StatusBadRequest, NewErrorResp(http.StatusBadRequest, "Session can not be set when open a stateful connection"))
	}
	if action == StatefulActionClose && !strings.HasPrefix(session, StatefulNonTransactionSessionPrefix) {
		ctx.JSON(http.StatusBadRequest, NewErrorResp(http.StatusBadRequest, "Stateful session must be set when close a stateful connection"))
	}
	// check request
	var req ExecuteRequest
	if err := ctx.ShouldBindJSON(&req); err != nil {
		logger.Warn("bind json error", zap.Error(err))
		ctx.JSON(http.StatusBadRequest, NewErrorResp(http.StatusBadRequest, InvalidBodyErrorMessage))
		return
	}
	logger.Debug("query", zap.String("query", req.Query))
	// check query
	sqlType, err1 := parseSqlType(ctx, req.Query)
	if err1 != nil {
		logger.Error("failed to parse sql type", zap.Error(err1))
		ctx.JSON(http.StatusBadRequest, NewErrorResp(http.StatusBadRequest, fmt.Sprintf("%s: %s", ParseSQLFailMessage, err1.Error())))
		return
	}
	// get session type
	sessionType := GetSessionType(session, sqlType, ctx.GetHeader(StatefulActionHeader))

	// compatible code, delete it after it has been released
	// if the session is stateful non transaction session, set it to stateless session and set session to empty
	if sessionType == UnknownSession {
		logger.Warn("session hit compatible code", zap.String("session", session))
		sessionType = StatelessSession
		session = ""
	}

	var resp *ExecuteResponse
	var err bizerr.BizError
	switch sessionType {
	case StatelessSession:
		resp, err = g.HandleStatelessRequest(ctx, session, req, sqlType)
	case StatefulTransactionSession:
		resp, err = g.HandleStatefulTxnRequest(ctx, session, req, sqlType)
	case StatefulNonTransactionSession:
		resp, err = g.HandleStatefulNonTxnRequest(ctx, session, req, sqlType)
	default:
		logger.Error("unknown session type", zap.String("sessionType", string(sessionType)))
		err = bizerr.New(bizerr.ErrCodeParams, InvalidSessionErrorMessage)
	}

	if err != nil {
		ctx.JSON(err.HttpStatus(), BizErrorToErrorResp(err))
	} else {
		ctx.JSON(http.StatusOK, resp)
	}
}

func (g *HttpGateway) HandleStatelessRequest(ctx *gin.Context, session string, req ExecuteRequest, sqlType SqlType) (resp *ExecuteResponse, err bizerr.BizError) {
	ctxReq := ctx.Request.Context()
	logger := log.WithContextL(ctxReq).With(zap.String("session", session))

	// get or create db pool
	var pool *DBPool
	if session != "" {
		pool, err = g.poolManager.GetOrCreateByToken(ctxReq, GetStatelessToken(session))
		if err != nil {
			logger.Error("failed to get db from pool", zap.Error(err))
			return nil, err
		}
	} else {
		cred, err := getCredFromCtx(ctx)
		if err != nil {
			logger.Error("failed to get credential from context", zap.Error(err))
			return nil, err
		}
		pool, err = g.poolManager.GetOrCreateByCred(ctxReq, cred)
		if err != nil {
			logger.Error("failed to get db from pool", zap.Error(err))
			return nil, err
		}
	}
	resp, err = pool.Query(ctxReq, req.Query, sqlType)
	if err != nil {
		logger.Error("failed to execute query", zap.Error(err))
		return nil, err
	}
	ctx.Header(SessionHeader, GenerateStatelessSession(pool.token))
	return resp, nil
}

func (g *HttpGateway) HandleStatefulTxnRequest(ctx *gin.Context, session string, req ExecuteRequest, sqlType SqlType) (resp *ExecuteResponse, err bizerr.BizError) {
	ctxReq := ctx.Request.Context()
	logger := log.WithContextL(ctxReq).With(zap.String("session", session))

	// check session and query
	if session != "" && sqlType == Begin {
		logger.Warn("txn session can not be set when begin a transaction")
		return nil, bizerr.New(bizerr.ErrCodeParams, BeginWithSessionMessage)
	}
	if session == "" && !(sqlType == Begin) {
		logger.Warn("txn session must be set when execute transaction")
		return nil, bizerr.New(ErrCodeEmptyTx, TxnWithoutSessionMessage)
	}

	// get or create txn session
	var txnSession *TxnSession
	if session != "" {
		txnSession = g.sessionStore.GetTxn(session)
		if txnSession == nil {
			logger.Error("txn session not found")
			return nil, bizerr.New(ErrCodeEmptyTx, TxNotFoundErrorMessage)
		}
	} else {
		isolationLevel := ctx.GetHeader(TxnIsolationLevelHeader)
		cred, err := getCredFromCtx(ctx)
		if err != nil {
			logger.Error("failed to get credential from context", zap.Error(err))
			return nil, err
		}
		pool, err := g.poolManager.GetOrCreateByCred(ctxReq, cred)
		if err != nil {
			logger.Error("failed to get db from pool", zap.Error(err))
			return nil, err
		}
		// use background context to avoid tx rollback dut to ctx cancelation
		txCtx := context.Background()
		tx, err := pool.BeginTx(txCtx, isolationLevel)
		if err != nil {
			logger.Error("failed to begin transaction", zap.Error(err))
			return nil, err
		}
		txnSession = g.sessionStore.NewTxnSession(txCtx, tx, cred, g.order)
		logger.Info("begin tx success", zap.String("session", txnSession.session), zap.String("username", cred.Username))
		ctx.Header(SessionHeader, txnSession.session)
		return &ExecuteResponse{}, nil
	}

	// execute query
	resp, err = txnSession.Query(ctxReq, req.Query, sqlType)
	if err != nil {
		logger.Error("failed to execute query", zap.Error(err))
		return nil, err
	}
	ctx.Header(SessionHeader, txnSession.session)
	return resp, nil
}

func (g *HttpGateway) HandleStatefulNonTxnRequest(ctx *gin.Context, session string, req ExecuteRequest, sqlType SqlType) (resp *ExecuteResponse, err bizerr.BizError) {
	ctxReq := ctx.Request.Context()
	logger := log.WithContextL(ctxReq).With(zap.String("session", session))

	// handle close action first
	action := strings.ToLower(strings.TrimSpace(ctx.GetHeader(StatefulActionHeader)))
	if action == StatefulActionClose {
		if session == "" {
			logger.Error("close non-txn session without session")
			return nil, bizerr.New(bizerr.ErrCodeParams, InvalidSessionErrorMessage)
		}
		nonTxnSession := g.sessionStore.GetNonTxn(session)
		if nonTxnSession == nil {
			logger.Warn("non-txn session not found, ignore close")
			return &ExecuteResponse{}, nil
		}
		if err := nonTxnSession.Close(ctxReq); err != nil {
			logger.Error("failed to close non-txn session", zap.Error(err))
			return nil, err
		}
		ctx.Header(SessionHeader, nonTxnSession.session)
		return &ExecuteResponse{}, nil
	}

	// get or create non-txn session
	var nonTxnSession *NonTxnSession
	if session != "" {
		nonTxnSession = g.sessionStore.GetNonTxn(session)
		if nonTxnSession == nil {
			logger.Error("non-txn session not found")
			return nil, bizerr.New(bizerr.ErrCodeParams, StatefulNotFoundErrorMessage)
		}
	} else {
		cred, err := getCredFromCtx(ctx)
		if err != nil {
			logger.Error("failed to get credential from context", zap.Error(err))
			return nil, err
		}
		pool, err := g.poolManager.GetOrCreateByCred(ctxReq, cred)
		if err != nil {
			logger.Error("failed to get db from pool", zap.Error(err))
			return nil, err
		}
		connCtx := context.Background()
		conn, err := pool.Connect(connCtx)
		if err != nil {
			logger.Error("failed to connect db", zap.Error(err))
			return nil, err
		}
		nonTxnSession = g.sessionStore.NewNonTxnSession(connCtx, conn, cred, g.order)
		logger.Info("create non-txn session success", zap.String("session", nonTxnSession.session), zap.String("username", cred.Username))
		if req.Query == "" {
			ctx.Header(SessionHeader, nonTxnSession.session)
			return &ExecuteResponse{}, nil
		}
	}

	// execute query
	resp, err = nonTxnSession.Query(ctxReq, req.Query, sqlType)
	if err != nil {
		logger.Error("failed to execute query", zap.Error(err))
		return nil, err
	}
	ctx.Header(SessionHeader, nonTxnSession.session)
	return resp, nil
}

func getCredFromCtx(ctx *gin.Context) (*Credential, bizerr.BizError) {
	authHeader := ctx.Request.Header.Get("Authorization")
	if len(authHeader) == 0 {
		return nil, bizerr.New(bizerr.ErrCodeAuth, EmptyAuthErrorMessage)
	}
	username, password := extractBasicAuth(authHeader)
	if username == "" {
		return nil, bizerr.New(bizerr.ErrCodeAuth, InvalidAuthErrorMessage+": missing username")
	}
	if password == "" {
		return nil, bizerr.New(bizerr.ErrCodeAuth, InvalidAuthErrorMessage+": missing password")
	}
	database := ctx.GetHeader(DatabaseHeader)
	return &Credential{
		Username: username,
		Password: password,
		Database: database,
	}, nil
}

func extractBasicAuth(authHeader string) (string, string) {
	basicAuthPrefix := "Basic"
	authHeaderLen := 2
	auth := strings.SplitN(authHeader, " ", authHeaderLen)
	if len(auth) != authHeaderLen || auth[0] != basicAuthPrefix {
		return "", ""
	}

	payload, err := base64.StdEncoding.DecodeString(auth[1])
	if err != nil {
		return "", ""
	}
	basicAuthLen := 2
	pair := strings.SplitN(string(payload), ":", basicAuthLen)
	if len(pair) != basicAuthLen {
		return "", ""
	}

	return pair[0], pair[1]
}

func (g *HttpGateway) forward(ctx *gin.Context, order int) error {
	logger := log.WithContextL(ctx.Request.Context())
	source := getPodHost(g.order, g.componentName)
	target := getPodHost(order, g.componentName)
	logger.Info(fmt.Sprintf("forward request from %s to %s", source, target))

	director := func(req *http.Request) {
		req.URL.Scheme = "https"
		req.URL.Host = target
		req.URL.Path = ctx.Request.URL.Path + "/proxy"
	}
	proxy := &httputil.ReverseProxy{Director: director}
	// no need to verify cert inside k8s cluster
	//nolint:mnd,gosec // ignore
	proxy.Transport = &http.Transport{
		IdleConnTimeout:     60 * time.Second,
		TLSHandshakeTimeout: 5 * time.Second,
		TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
		ForceAttemptHTTP2:   true,
	}
	// Inherit httpproxy.defaulterrorhandle but change the log to output logid/traceid
	proxy.ErrorHandler = func(writer http.ResponseWriter, request *http.Request, e error) {
		logger.Error("http: proxy error", zap.Error(e))
		writer.WriteHeader(http.StatusBadGateway)
	}

	// ctx.Request.Clone() does not clone the body, and body will be close in ServerHTTP.
	// So we use NopCloser to replace it.
	ByteBody, err := io.ReadAll(ctx.Request.Body)
	if err != nil {
		logger.Error("read request body failed", zap.Error(err))
		return err
	}
	ctx.Request.Body = io.NopCloser(bytes.NewBuffer(ByteBody))
	proxy.ServeHTTP(ctx.Writer, ctx.Request)
	if ctx.Writer.Status() == http.StatusBadGateway {
		return errors.New("proxy with 502")
	}
	return nil
}

func getPodHost(order int, componentName string) string {
	return fmt.Sprintf("%s-%d.%s", componentName, order, componentName)
}