package httpgateway

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

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

type NonTxnSession struct {
	session string
	conn    *sql.Conn
	cred    *Credential
	expire  time.Time
	inUse   bool
	mu      sync.Mutex
	store   *StatefulSessionStore
	connCtx context.Context
}

func (s *StatefulSessionStore) NewNonTxnSession(connCtx context.Context, conn *sql.Conn, cred *Credential, order int) *NonTxnSession {
	session := GenerateStatefulNonTransactionSession(order)
	nonTxnSession := &NonTxnSession{
		session: session,
		conn:    conn,
		cred:    cred,
		expire:  time.Now().Add(StatefulTTL),
		inUse:   false,
		store:   s,
		connCtx: connCtx,
	}
	s.nonTxnSessions.Store(session, nonTxnSession)
	return nonTxnSession
}

func (s *StatefulSessionStore) GetNonTxn(session string) *NonTxnSession {
	value, ok := s.nonTxnSessions.Load(session)
	if !ok {
		return nil
	}
	if value == nil {
		return nil
	}
	return value.(*NonTxnSession)
}

func (s *StatefulSessionStore) DeleteNonTxn(ctx context.Context, session string) {
	logger := log.WithContextL(ctx)
	logger.Info("clean non-txn session", zap.String("session", session))
	s.nonTxnSessions.Delete(session)
}

func (s *StatefulSessionStore) RangeNonTxn(fn func(string, *NonTxnSession) bool) {
	s.nonTxnSessions.Range(func(key, value any) bool {
		return fn(key.(string), value.(*NonTxnSession))
	})
}

func (s *NonTxnSession) tryToUseNonTxn() bizerr.BizError {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.inUse {
		return bizerr.New(bizerr.ErrCodeParams, StatefulConcurrencyErrorMessage)
	}
	if time.Now().After(s.expire) {
		return bizerr.New(ErrCodeSessionExpired, StatefulExpiredErrorMessage)
	}
	s.inUse = true
	return nil
}

func (s *NonTxnSession) tryToCleanNonTxn(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.conn.Close(); err != nil {
			logger.Error("failed to close non-txn session", zap.String("session", s.session), zap.Error(err))
			if !strings.Contains(err.Error(), "connection is already closed") {
				return false
			}
		}
		s.store.DeleteNonTxn(ctx, s.session)
		return true
	}
	return false
}

func (s *NonTxnSession) 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(StatefulNonTransactionSession)))

	if err := s.tryToUseNonTxn(); err != nil {
		logger.Error("try to use non-txn 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 Query:
		rows, err := s.conn.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.conn.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.conn.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)
	}
}

func (s *NonTxnSession) Close(ctx context.Context) bizerr.BizError {
	logger := log.WithContextL(ctx).With(zap.String("session", s.session)).With(zap.String("username", s.cred.Username))
	logger.Info("start to close", zap.String("session type", string(StatefulNonTransactionSession)))

	// close is concurrent safe and will be blocked until all other operations finish
	if err := s.conn.Close(); err != nil {
		logger.Error("failed to close non-txn session", zap.Error(err))
		return bizerr.New(ErrCodeQueryFail, "Close connection failed: %s", err.Error())
	}
	s.store.DeleteNonTxn(ctx, s.session)
	return nil
}