package httpgateway

import (
	"context"
	"database/sql"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/lithammer/shortuuid/v4"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/log"
	"go.uber.org/zap"
)

type SessionType string

const (
	StatefulTransactionSession    SessionType = "StatefulTransaction"
	StatefulNonTransactionSession SessionType = "StatefulNonTransaction"
	StatelessSession              SessionType = "Stateless"
	UnknownSession                SessionType = "Unknown"
)

const (
	StatefulTransactionSessionPrefix    = "txn_"
	StatefulNonTransactionSessionPrefix = "stateful_"
	StatelessSessionPrefix              = "stateless_"
	StatefulTTL                         = 10 * time.Minute
)

type TxnSession struct {
	session string
	tx      *sql.Tx
	cred    *Credential
	expire  time.Time
	inUse   bool
	mu      sync.Mutex
	store   *StatefulSessionStore
	txCtx   context.Context
}

type StatefulSessionStore struct {
	// session(string) -> TxnSession
	txnSessions sync.Map
	// session(string) -> NonTxnSession
	nonTxnSessions sync.Map
}

func NewStatefulSessionStore() *StatefulSessionStore {
	return &StatefulSessionStore{}
}

func (s *StatefulSessionStore) NewTxnSession(txCtx context.Context, tx *sql.Tx, cred *Credential, order int) *TxnSession {
	session := GenerateStatefulTransactionSession(order)
	txnSession := &TxnSession{
		session: session,
		tx:      tx,
		cred:    cred,
		expire:  time.Now().Add(StatefulTTL),
		inUse:   false,
		store:   s,
		txCtx:   txCtx,
	}
	s.txnSessions.Store(session, txnSession)
	return txnSession
}

func (s *StatefulSessionStore) GetTxn(session string) *TxnSession {
	value, ok := s.txnSessions.Load(session)
	if !ok {
		return nil
	}
	if value == nil {
		return nil
	}
	return value.(*TxnSession)
}

func (s *StatefulSessionStore) DeleteTxn(ctx context.Context, session string) {
	logger := log.WithContextL(ctx)
	logger.Info("clean session", zap.String("session", session))
	s.txnSessions.Delete(session)
}

func (s *StatefulSessionStore) RangeTxn(fn func(string, *TxnSession) bool) {
	s.txnSessions.Range(func(key, value any) bool {
		return fn(key.(string), value.(*TxnSession))
	})
}

func GetStatefulSessionOrder(session string) (int, error) {
	if session == "" {
		return 0, fmt.Errorf("empty session: %s", session)
	}
	splits := strings.Split(session, "_")
	order, err := strconv.Atoi(splits[len(splits)-1])
	if err != nil {
		return 0, fmt.Errorf("invalid session: %s", session)
	}
	return order, nil
}

func GetSessionType(session string, sqlType SqlType, statefulHeaderValue string) SessionType {
	if session == "" {
		if sqlType == Begin {
			return StatefulTransactionSession
		}
		if strings.ToLower(statefulHeaderValue) == StatefulActionOpen {
			return StatefulNonTransactionSession
		}
		return StatelessSession
	}
	if strings.HasPrefix(session, StatelessSessionPrefix) {
		return StatelessSession
	}
	if strings.HasPrefix(session, StatefulTransactionSessionPrefix) {
		return StatefulTransactionSession
	}
	if strings.HasPrefix(session, StatefulNonTransactionSessionPrefix) {
		return StatefulNonTransactionSession
	}
	return UnknownSession
}

func GetStatelessToken(session string) string {
	return strings.TrimPrefix(session, StatelessSessionPrefix)
}

func GenerateStatelessSession(token string) string {
	return fmt.Sprintf("%s%s", StatelessSessionPrefix, token)
}

func GenerateStatefulTransactionSession(order int) string {
	return fmt.Sprintf("%s%s_%d", StatefulTransactionSessionPrefix, shortuuid.New(), order)
}

func GenerateStatefulNonTransactionSession(order int) string {
	return fmt.Sprintf("%s%s_%d", StatefulNonTransactionSessionPrefix, shortuuid.New(), order)
}

func (s *TxnSession) tryToUseTxn() bizerr.BizError {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.inUse {
		return bizerr.New(bizerr.ErrCodeParams, TxConcurrencyErrorMessage)
	}
	if time.Now().After(s.expire) {
		return bizerr.New(ErrCodeSessionExpired, TxExpiredErrorMessage)
	}
	s.inUse = true
	return nil
}

func (s *TxnSession) tryToCleanTxn(ctx context.Context) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	logger := log.WithContextL(ctx)
	if time.Now().After(s.expire) && !s.inUse {
		if err := s.tx.Rollback(); err != nil {
			logger.Error("failed to rollback txn session", zap.String("session", s.session), zap.Error(err))
			if !strings.Contains(err.Error(), "transaction has already been committed or rolled back") {
				return false
			}
		}
		s.store.DeleteTxn(ctx, s.session)
		return true
	}
	return false
}

func (s *TxnSession) Query(ctx context.Context, query string, sqlType SqlType) (*ExecuteResponse, bizerr.BizError) {
	logger := log.WithContextL(ctx).With(zap.String("session", s.session)).With(zap.String("username", s.cred.Username))
	logger.Info("start to query", zap.String("session type", string(StatefulTransactionSession)))

	if err := s.tryToUseTxn(); err != nil {
		logger.Error("try to use transaction session failed", zap.Error(err))
		return nil, err
	}
	defer func() {
		s.mu.Lock()
		s.expire = time.Now().Add(StatefulTTL)
		s.inUse = false
		s.mu.Unlock()
	}()

	switch sqlType {
	case Commit:
		err := s.tx.Commit()
		if err != nil {
			logger.Error("failed to commit transaction", zap.Error(err))
			return nil, bizerr.New(ErrCodeCommitFail, ErrCodeCommitFailMessage+": %s", err.Error())
		}
		s.store.DeleteTxn(ctx, s.session)
		return &ExecuteResponse{}, nil
	case Rollback:
		err := s.tx.Rollback()
		if err != nil {
			logger.Error("failed to rollback transaction", zap.Error(err))
			return nil, bizerr.New(ErrCodeRollBackFail, ErrCodeRollBackFailMessage+": %s", err.Error())
		}
		s.store.DeleteTxn(ctx, s.session)
		return &ExecuteResponse{}, nil
	case Query:
		rows, err := s.tx.QueryContext(ctx, query)
		if err != nil {
			logger.Error("failed to query db", zap.Error(err))
			return nil, bizerr.New(ErrCodeQueryFail, ErrCodeQueryFailMessage+": %s", err.Error())
		}
		defer func(rows *sql.Rows) {
			err := rows.Close()
			if err != nil {
				logger.Error("failed to close rows", zap.Error(err))
			}
		}(rows)
		resp, err := parseRows(ctx, rows)
		if err != nil {
			logger.Error("failed to parse rows", zap.Error(err))
			return nil, bizerr.New(ErrCodeParseRowFail, UnknownInternalErrorMessage)
		}
		return resp, nil
	case DML:
		result, err := s.tx.ExecContext(ctx, query)
		if err != nil {
			logger.Error("failed to exec sql", zap.Error(err))
			return nil, bizerr.New(ErrCodeQueryFail, ErrCodeQueryFailMessage+": %s", err.Error())
		}
		resp, err := parseResult(ctx, result)
		if err != nil {
			logger.Error("failed to parse result", zap.Error(err))
			return nil, bizerr.New(ErrCodeParseResultFail, UnknownInternalErrorMessage)
		}
		return resp, nil
	case Execute:
		_, err := s.tx.ExecContext(ctx, query)
		if err != nil {
			logger.Error("failed to exec ddl", zap.Error(err))
			return nil, bizerr.New(ErrCodeQueryFail, ErrCodeQueryFailMessage+": %s", err.Error())
		}
		return &ExecuteResponse{}, nil
	default:
		logger.Error("Unsupported SQL")
		return nil, bizerr.New(bizerr.ErrCodeParams, UnsupportedSQLMessage)
	}
}