package sessioncache

import (
	"context"
	"database/sql"
	"dataservices/commons/config"
	"dataservices/commons/errs"
	"dataservices/commons/logger"
	"dataservices/commons/metrics"
	"dataservices/dbhubsdk"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
)

type SessionType int8

const (
	SessionType_Standard SessionType = iota
	SessionType_Once
	SessionType_Endpoint
)

type ConnectionBuilder func(ctx context.Context, database string, cmds []Cmd) (*sql.DB, error)

type SessionItem struct {
	SessionID        string      `json:"session_id"`
	EndSec           uint64      `json:"end_sec"`
	Database         string      `json:"database"`
	LastActiveTime   uint64      `json:"last_active_time"`
	CrtID            string      `json:"crt_id"`
	ClusterID        uint64      `json:"cluster_id"`
	ExpireSec        uint64      `json:"expire_sec"`
	Cmds             []Cmd       `json:"cmds"`              // The status command buffer records operations related to the session status, such as set @ a=1;
	ExecutorEndpoint string      `json:"executor_endpoint"` // Record the bound executor pod IP address
	Type             SessionType `json:"type"`

	dbPool      dbhubsdk.DBConnPool `json:"-"`
	conn        *sql.DB             `json:"-"`
	reconnector ConnectionBuilder   `json:"-"`
	modifyFlag  bool                `json:"-"`
	mutex       sync.Mutex
}

type Cmd struct {
	Type    string `json:"type"`
	Content string `json:"content"`
}

// Database configuration
type SessionCache struct {
	maxNum     uint64
	sessionMap sync.Map
}

var storage = &SessionCache{}

func closeSessionDB(sessionItem *SessionItem) {
	if sessionItem == nil || sessionItem.Type != SessionType_Standard || sessionItem.conn == nil {
		return
	}
	_ = sessionItem.conn.Close()
}

func clearExpiredSessions(storage *SessionCache, curTimestamp int64) int {
	num := 0
	storage.sessionMap.Range(func(k, v interface{}) bool {
		val := v.(*SessionItem)
		if curTimestamp > int64(val.EndSec) {
			storage.sessionMap.Delete(k)
			closeSessionDB(val)
			tilog.Info(context.Background(), "cronClearExpireSession Close Session", zap.String("session_id", k.(string)))
		}
		num += 1
		return true
	})
	return num
}

func cronClearExpireSession(storage *SessionCache) {
	for {
		num := clearExpiredSessions(storage, time.Now().Unix())
		metrics.SessionConnTotal.With(prometheus.Labels{"session": "total"}).Set(float64(num))
		time.Sleep(time.Millisecond * 100)
	}
}

func CloseSession(sessionId string) error {
	sessionItem, err := loadSession(sessionId)
	if err != nil {
		return err
	}
	sessionItem.EndSec = uint64(time.Now().Unix())
	return nil
}

func RefreshSession(sessionId string) error {
	sessionItem, err := loadSession(sessionId)
	if err != nil {
		return err
	}
	sessionItem.modifyFlag = true
	sessionItem.LastActiveTime = uint64(time.Now().Unix())
	sessionItem.EndSec = uint64(time.Now().Unix()) + sessionItem.ExpireSec
	return nil
}

func GetSession(sessionId string) (*SessionItem, error) {
	sessionItem, err := loadSession(sessionId)
	if err != nil {
		return nil, err
	}
	return sessionItem, nil
}

func CreateOnceSession(dbPool dbhubsdk.DBConnPool, clusterID uint64) *SessionItem {
	return &SessionItem{
		Type:      SessionType_Once,
		dbPool:    dbPool,
		ClusterID: clusterID,
	}
}

func CreateEndpointSession(sessionID string, dbPool dbhubsdk.DBConnPool, clusterID uint64) *SessionItem {
	return &SessionItem{
		Type:             SessionType_Endpoint,
		SessionID:        sessionID,
		dbPool:           dbPool,
		ClusterID:        clusterID,
		ExpireSec:        300,
		ExecutorEndpoint: config.RpcEndpointAddr(),
		modifyFlag:       true,
	}
}

func CreateSession(sessionID string, sqlDB *sql.DB, database, crtID string, clusterID, expireSec uint64, cmds []Cmd, reconnector ConnectionBuilder) *SessionItem {
	if sessionID == "" {
		sessionID = uuid.NewString()
	}
	item := &SessionItem{
		Type:             SessionType_Standard,
		SessionID:        sessionID,
		conn:             sqlDB,
		EndSec:           uint64(time.Now().Unix()) + expireSec,
		Database:         database,
		reconnector:      reconnector,
		Cmds:             cmds,
		LastActiveTime:   uint64(time.Now().Unix()),
		CrtID:            crtID,
		ClusterID:        clusterID,
		ExpireSec:        expireSec,
		ExecutorEndpoint: config.RpcEndpointAddr(),
		modifyFlag:       true,
	}
	storage.sessionMap.Store(sessionID, item)
	return item
}

func ResetConn(sessionId string, sqlDB *sql.DB) error {
	sessionItem, err := loadSession(sessionId)
	if err != nil {
		return err
	}
	sessionItem.mutex.Lock()
	defer sessionItem.mutex.Unlock()

	oldConn := sessionItem.conn
	sessionItem.modifyFlag = true
	sessionItem.conn = sqlDB
	sessionItem.EndSec = uint64(time.Now().Unix()) + sessionItem.ExpireSec
	if oldConn != nil && oldConn != sqlDB {
		_ = oldConn.Close()
	}
	return nil
}

func ResetDBName(sessionId string, dbName string) error {
	sessionItem, err := loadSession(sessionId)
	if err != nil {
		return err
	}
	sessionItem.mutex.Lock()
	defer sessionItem.mutex.Unlock()

	sessionItem.modifyFlag = true
	sessionItem.Database = dbName
	sessionItem.EndSec = uint64(time.Now().Unix()) + sessionItem.ExpireSec
	return nil
}

func Init(maxNum uint64) error {
	storage.maxNum = maxNum
	go cronClearExpireSession(storage)
	return nil
}

func loadSession(sessionId string) (*SessionItem, error) {
	val, ok := storage.sessionMap.Load(sessionId)
	if !ok {
		return nil, errors.New("no sessionId found in cache: " + sessionId)
	}
	return val.(*SessionItem), nil
}

func (s *SessionItem) AppendCmd(cmd Cmd) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	s.Cmds = append(s.Cmds, cmd)
	n := len(s.Cmds)
	cmdExists := make(map[Cmd]bool, n)
	idx := n - 1
	for i := n - 1; i >= 0; i-- {
		if cmdExists[s.Cmds[i]] {
			continue
		}
		cmdExists[s.Cmds[i]] = true
		if idx != i {
			s.Cmds[idx] = s.Cmds[i]
		}
		idx--
	}
	s.Cmds = s.Cmds[idx+1:]

	if len(s.Cmds) > 20 {
		s.Cmds = s.Cmds[len(s.Cmds)-20:]
	}
	s.modifyFlag = true
}

func (s *SessionItem) RequestPost(ctx context.Context) error {
	if s.Type == SessionType_Once {
		return nil
	}
	return s.Save(ctx)
}

func (s *SessionItem) Save(ctx context.Context) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.Type == SessionType_Once {
		return nil
	}

	if !s.modifyFlag {
		return nil
	}

	var (
		begin time.Time
		ok    string
		err   error
	)
	begin = time.Now()
	err = driver().Save(ctx, s)
	if err != nil {
		ok = "false"
		tilog.Warn(ctx, "save session to remote db failed", zap.String("session_id", s.SessionID), zap.Error(err))
	} else {
		ok = "true"
		s.modifyFlag = false
	}
	metrics.SessionCacheDuration.WithLabelValues("save", ok).Observe(float64(time.Since(begin).Milliseconds()))
	metrics.SessionCacheCounter.WithLabelValues("save", ok).Inc()
	return err
}

func (s *SessionItem) Delete(ctx context.Context) error {
	var (
		begin time.Time
		ok    string
		err   error
	)
	begin = time.Now()
	err = driver().Delete(ctx, s.SessionID)
	if err != nil {
		ok = "false"
		tilog.Warn(ctx, "delete session failed", zap.String("session_id", s.SessionID), zap.Error(err))
	} else {
		ok = "true"
	}
	metrics.SessionCacheDuration.WithLabelValues("delete", ok).Observe(float64(time.Since(begin).Milliseconds()))
	metrics.SessionCacheCounter.WithLabelValues("delete", ok).Inc()
	return err
}

func (s *SessionItem) Load(ctx context.Context) error {
	var (
		begin time.Time
		ok    string
		err   error
	)
	begin = time.Now()
	err = driver().Load(ctx, s)
	if err != nil {
		ok = "false"
		tilog.Warn(ctx, "load session failed", zap.String("session_id", s.SessionID), zap.Error(err))
	} else {
		ok = "true"
	}
	metrics.SessionCacheDuration.WithLabelValues("load", ok).Observe(float64(time.Since(begin).Milliseconds()))
	metrics.SessionCacheCounter.WithLabelValues("load", ok).Inc()
	return err
}

func (s *SessionItem) IsStandard() bool {
	return s.Type == SessionType_Standard
}

func (s *SessionItem) Handle(ctx context.Context, times int, handler func(context.Context, *sql.Conn) (bool, error)) error {
	switch s.Type {
	case SessionType_Standard:
		if times > 0 {
			db, rerr := s.reconnector(ctx, s.Database, s.Cmds)
			if rerr != nil {
				return errs.NewInternalDB("Reconnect db failed: "+s.SessionID, rerr)
			}
			ResetConn(s.SessionID, db)
		}
		startTime := time.Now()
		conn, err := s.conn.Conn(ctx)
		logger.Info(ctx, "standard sesssion handler get connection complete", zap.Duration("cost", time.Since(startTime)))
		if err != nil {
			return err
		}
		defer conn.Close()
		_, err = handler(ctx, conn)
		return err

	case SessionType_Once, SessionType_Endpoint:
		return s.dbPool.Handle(ctx, handler)
	default:
		return errs.NewBadRequest(fmt.Sprintf("unsupported session type:%d", s.Type), nil)
	}
}

func GenEndpointSessionID(endpointID uint64) string {
	return fmt.Sprintf("dataservice_endpoint_%d", endpointID)
}

func IsEndpointSessionID(sessionID string) bool {
	return strings.HasPrefix(sessionID, "dataservice_endpoint_")
}