// Package dbhubsdk add by blacknc <blacknc@163.com> at 2023/11/02 17:45:06
package dbhubsdk

import (
	"context"
	"database/sql"
	"dataservices/commons/config"
	"dataservices/commons/contextvar"
	"dataservices/commons/errs"
	"dataservices/commons/logger"
	"dataservices/commons/metrics"
	"dataservices/commons/utils"
	"dataservices/internal/entity"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
)

const (
	DefaultMaxIdelTime time.Duration = 5 * time.Minute
	DefaultMaxOpenConn int           = 50
	DefaultMaxIdelConn int           = 20
)

var pool *dbHubPool

var ErrPoolClosed = errs.NewInternalDB("db connecton pool closed", nil)

func InitHubPool(ctx context.Context) {
	pool = &dbHubPool{ctx: ctx}
	runDBPoolCleaner(pool)
}

func HubPool() DBHubPool {
	return pool
}

func GetConnPool(ctx context.Context, in *entity.ConnectionInput) (DBConnPool, error) {
	return pool.GetConnPool(ctx, in)
}

type DBHubPool interface {
	GetConnPool(ctx context.Context, in *entity.ConnectionInput) (DBConnPool, error)
}

type DBConnPool interface {
	GetConn(context.Context) (*sql.Conn, error)
	Handle(context.Context, func(context.Context, *sql.Conn) (bool, error)) error
}

type refreshRequest struct {
	notifych chan struct{}
}

type dbConnPool struct {
	dbs          [2]*sql.DB
	idx          atomic.Int32
	expireAt     atomic.Int64
	jwtMd5       atomic.Value
	lastActive   atomic.Int64
	lastRefresh  atomic.Int64
	createdAt    int64
	connectionIn *entity.ConnectionInput
	poolKey      string
	closed       atomic.Bool
	closech      chan struct{}
	readych      chan struct{}
	refreshch    chan *refreshRequest
	err          error
}

func (p *dbConnPool) getLogFields() []zap.Field {
	fields := []zap.Field{
		zap.String("poolKey", p.poolKey),
		zap.Int64("poolCreatedAt", p.createdAt),
		zap.Uint64("poolProject", p.connectionIn.ProjectId),
		zap.Uint64("poolCluster", p.connectionIn.ClusterId),
		zap.Any("poolJwtMd5WithSalt", p.jwtMd5.Load()),
		zap.Int64("poolJwtExpireAt", p.expireAt.Load()),
		zap.Int64("poolLastActive", p.lastActive.Load()),
		zap.Int64("poolLastRefresh", p.lastRefresh.Load()),
		zap.String("poolCretType", p.connectionIn.Credential.GetType().String()),
		logger.Redact("poolCretAccount", p.connectionIn.Credential.GetAccountId()),
	}
	fields = append(fields, zap.String("poolCretKeyMd5WithSalt", utils.Md5WithDefaultSalt(p.connectionIn.Credential.GetKey())))
	return fields
}

// GetConn Get one database connection from the connection pool,
// Don't forget to release the connection after use!!!
func (p *dbConnPool) GetConn(ctx context.Context) (conn *sql.Conn, err error) {
	startTime := time.Now()
	logger := contextvar.Logger(ctx).With(p.getLogFields()...).With(zap.Time("startTime", startTime))
	logger.Info("DBConnPool.GetConn begin")
	defer func() {
		logger.Info("DBConnPool.GetConn completed", zap.Duration("cost", time.Since(startTime)), zap.Error(err))
	}()

	if p.closed.Load() {
		return nil, ErrPoolClosed
	}
	p.lastActive.Store(time.Now().Unix())

	conn, err = p.dbs[p.idx.Load()].Conn(ctx)
	if err == nil {
		return
	}

	logger.Info("DBConnPool.GetConn get connection failed, Notify the updater to refresh connection and then try again", zap.Error(err))

	// refresh db pool
	err = nil
	notifych := make(chan struct{})
	select {
	case <-ctx.Done():
		err = errs.NewInternalDB("request refresh connection of db pool timeout", ctx.Err())
	case p.refreshch <- &refreshRequest{notifych: notifych}:
		select {
		case <-ctx.Done():
			err = errs.NewInternalDB("wait refresh connection of db pool timeout", ctx.Err())
		case <-notifych:
		}
	}
	if err != nil {
		return
	}

	conn, err = p.dbs[p.idx.Load()].Conn(ctx)
	return
}

// Handle Inject the database connection into the user's handler and automatically call it,
// The handler does not need to worry about obtaining and releasing the connection.
func (p *dbConnPool) Handle(ctx context.Context, handler func(context.Context, *sql.Conn) (bool, error)) (err error) {
	startTime := time.Now()
	logger := contextvar.Logger(ctx).With(p.getLogFields()...).With(zap.Time("startTime", startTime))
	logger.Info("DBConnPool.Handle begin")
	defer func() {
		logger.Info("DBConnPool.Handle completed", zap.Duration("cost", time.Since(startTime)), zap.Error(err))
	}()

	// get connection from db pool
	conn, err := p.GetConn(ctx)
	logger = logger.With(zap.Duration("connection_cost", time.Since(startTime)))
	if err != nil {
		return err
	}

	// reset session if necessary
	var resetSession bool
	defer func() {
		closeStart := time.Now()
		if resetSession {
			conn.Raw(func(driverConn interface{}) error {
				if dc, ok := driverConn.(interface {
					ResetConnection(ctx context.Context) error
				}); ok {
					return dc.ResetConnection(ctx)
				}
				return nil
			})
		}
		conn.Close()
		logger = logger.With(zap.Duration("close_cost", time.Since(closeStart)), zap.Bool("reset_session", resetSession))
	}()

	// run user handler
	handlerStart := time.Now()
	resetSession, err = handler(ctx, conn)
	logger = logger.With(zap.Duration("handler_cost", time.Since(handlerStart)))
	logger.Info("DBConnPool.Handle execute user handler completed", zap.Error(err))
	return err
}

func (p *dbConnPool) close(reason string) {
	if p.closed.Swap(true) {
		return
	}
	logger := tilog.WithContextL(context.TODO()).With(p.getLogFields()...)
	logger.Info("dbConnPool closed", zap.String("reason", reason))
	for _, db := range p.dbs {
		if db != nil {
			db.Close()
		}
	}
	close(p.closech)
}

type dbHubPool struct {
	ctx     context.Context
	hubPool sync.Map
	mu      sync.RWMutex
	closed  atomic.Bool
}

func (p *dbHubPool) GetConnPool(ctx context.Context, in *entity.ConnectionInput) (pool DBConnPool, err error) {
	logger := contextvar.Logger(ctx)
	startTime := time.Now()
	defer func() {
		logger.Info("DBHubPool.GetConnPool completed", zap.Time("startTime", startTime),
			zap.Uint64("cluster", in.ClusterId), zap.Duration("cost", time.Since(startTime)), zap.Error(err))
	}()
	if p.closed.Load() {
		return nil, ErrPoolClosed
	}

	maxIdleConn, maxOpenConn := DefaultMaxIdelConn, DefaultMaxOpenConn

	in.Force = true
	in.ExpireTime = DefaultMaxIdelTime
	in.MaxIdleConns = &maxIdleConn
	in.MaxOpenConns = &maxOpenConn
	poolKey := genPoolKey(in)

	poolVal, _ := p.hubPool.Load(poolKey)
	if connPool, ok := poolVal.(*dbConnPool); ok {
		<-connPool.readych
		return connPool, connPool.err
	}

	p.mu.Lock()
	poolVal, _ = p.hubPool.Load(poolKey)
	if connPool, ok := poolVal.(*dbConnPool); ok {
		p.mu.Unlock()
		<-connPool.readych
		return connPool, connPool.err
	}

	connectionIn := *in
	connectionIn.Database = ""
	connPool := &dbConnPool{
		readych:      make(chan struct{}),
		refreshch:    make(chan *refreshRequest, 100),
		connectionIn: &connectionIn,
		poolKey:      poolKey,
	}
	connPool.lastActive.Store(time.Now().Unix())
	connPool.lastRefresh.Store(time.Now().Unix())
	p.hubPool.Store(poolKey, connPool)
	p.mu.Unlock()

	connOut, err := Serverless.BuildConnection(ctx, &connectionIn)
	if err != nil {
		p.mu.Lock()
		p.hubPool.Delete(poolKey)
		p.mu.Unlock()

		connPool.err = err
		close(connPool.readych)
		return nil, err
	}

	connPool.dbs = [2]*sql.DB{connOut.DB, nil}
	connPool.idx.Store(0)
	connPool.expireAt.Store(connOut.ExpireAt)
	connPool.jwtMd5.Store(connOut.JwtMd5)
	connPool.lastActive.Store(time.Now().Unix())
	connPool.lastRefresh.Store(time.Now().Unix())
	connPool.createdAt = time.Now().Unix()
	connPool.closech = make(chan struct{})
	runDBPoolUpdater(p.ctx, connPool)
	close(connPool.readych)

	logger.With(connPool.getLogFields()...).Info("GetConnPool Serverless.BuildConnection succeed")
	return connPool, nil
}

func genPoolKey(in *entity.ConnectionInput) string {
	account := in.Credential.GetAccountId()
	if account == "" {
		account = in.Credential.GetKey()
	}
	podAddr := config.RpcEndpointAddr()
	return utils.Md5(fmt.Sprintf("%d:%s:%s", in.ClusterId, in.Credential.GetType().String(), account)) + "-" + utils.Md5(podAddr)
}

func runDBPoolCleaner(pool *dbHubPool) {
	cleanerName := fmt.Sprintf("runDBPoolCleaner_%s", time.Now().Format("20060102150405"))
	utils.Worker(cleanerName, func() error {
		ticker := time.NewTicker(time.Minute)
		for {
			select {
			case <-pool.ctx.Done():
				tilog.Info(pool.ctx, "begin clean db pool because context done")
				if pool.closed.Swap(true) {
					pool.hubPool.Range(func(key, value any) bool {
						connPool, _ := value.(*dbConnPool)
						pool.hubPool.Delete(key)
						connPool.close("context done")
						return true
					})
				}
				return nil
			case <-ticker.C:
				tilog.Info(pool.ctx, "begin clean db pool because ticker")
				var currentAliveCount, refreshTimeoutCount, activeTimeoutCount int
				pool.hubPool.Range(func(key, value any) bool {
					connPool, _ := value.(*dbConnPool)
					if time.Now().Unix() > connPool.lastActive.Load()+300 {
						// No activity within 5 minutes, release connection pool
						pool.hubPool.Delete(key)
						connPool.close("max idle time")
						activeTimeoutCount++
					} else if time.Now().Unix() > connPool.lastRefresh.Load()+900 {
						// No updates within 15 minutes, release connection pool
						pool.hubPool.Delete(key)
						connPool.close("update timeout")
						refreshTimeoutCount++
					} else {
						currentAliveCount++
					}
					return true
				})
				metrics.DBHubPoolCounter.WithLabelValues("current_alive").Set(float64(currentAliveCount))
				metrics.DBHubPoolCounter.WithLabelValues("refresh_timeout").Set(float64(refreshTimeoutCount))
				metrics.DBHubPoolCounter.WithLabelValues("active_timeout").Set(float64(activeTimeoutCount))
			}
		}
	})
}

func runDBPoolUpdater(ctx context.Context, connPool *dbConnPool) {

	var (
		startTime         = time.Now()
		lastTryRefresh    = time.Now()
		interval          time.Duration
		pendingRefreshReq *refreshRequest
		maxInterval       = time.Minute * 10
		minInterval       = time.Millisecond * 100
	)

	updaterName := fmt.Sprintf("runDBPoolUpdater_%s_%s", connPool.poolKey, time.Now().Format("20060102150405"))
	loggerPrefix := tilog.WithContextL(ctx).With(zap.Time("startTime", startTime), zap.String("updaterName", updaterName))
	logger := func() *zap.Logger {
		return loggerPrefix.With(connPool.getLogFields()...)
	}
	utils.Worker(updaterName, func() error {

		for {
			// Notify the client that the connection has been refreshed
			if pendingRefreshReq != nil {
				close(pendingRefreshReq.notifych)
				pendingRefreshReq = nil
			}

			interval = time.Minute
			if expireAt := connPool.expireAt.Load(); expireAt > 0 {
				interval = time.Duration(expireAt-time.Now().Unix()-120) * time.Second
			}
			if interval < minInterval {
				interval = minInterval
			} else if interval > maxInterval {
				interval = maxInterval
			}

			select {
			case <-ctx.Done():
				logger().Info("context done")
				return nil

			case <-connPool.closech:
				logger().Info("connection pool closed")
				return nil

			case _refreshReq := <-connPool.refreshch:
				// No more repeated refreshes within 5 seconds to avoid frequent refreshes
				if time.Since(lastTryRefresh) < time.Second*5 {
					close(_refreshReq.notifych)
					continue
				}
				pendingRefreshReq = _refreshReq

			case <-time.After(interval):
			}

			logger().Info("begin update db pool")
			connOut, err := Serverless.BuildConnection(contextvar.WithLogger(ctx, logger()), connPool.connectionIn)
			if err != nil {
				logger().Warn("update db pool completed, Serverless.BuildConnection failed, waiting for next retry", zap.Error(err))
				lastTryRefresh = time.Now()
				continue
			}

			newIdx := 1 - connPool.idx.Load()
			if db := connPool.dbs[newIdx]; db != nil {
				if err := db.Close(); err != nil {
					logger().Warn("sql.DB close failed", zap.Error(err))
				}
			}
			connPool.dbs[newIdx] = connOut.DB
			connPool.idx.Store(newIdx)
			connPool.expireAt.Store(connOut.ExpireAt)
			connPool.jwtMd5.Store(connOut.JwtMd5)
			connPool.lastRefresh.Store(time.Now().Unix())
			lastTryRefresh = time.Now()
			logger().Info("update db pool completed, Serverless.BuildConnection succeed")
		}
	})
}