package httpgateway

import (
	"context"
	"crypto/tls"
	"database/sql"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/log"
	"go.uber.org/zap"
)

const (
	defaultDatabase = "test"
	defaultLimit    = "10000"

	poolMaxOpenConns = 1000
	poolMaxIdleConns = 1000
	poolConnMaxIdle  = 3 * time.Minute

	poolTTL = 1 * time.Hour

	TiDBGatewayService = "tidb-gateway-new.tidb-admin.svc"
)

var (
	mysqlSetupOnce sync.Once
)

type DBPool struct {
	db   *sql.DB
	cred *Credential
	// {username}-{password}-{database}
	sensitiveKey string
	expire       time.Time
	// token is the encrypted credential, may changed when db pool recreated
	token string
}

// PoolManager maintains sql.DB instances scoped by credential tuples.
type PoolManager struct {
	cipher      *SessionCipher
	gatewayHost string
	// sensitiveKey to DBPool
	pools sync.Map
	mu    sync.Mutex
}

func NewPoolManager(cipher *SessionCipher) *PoolManager {
	return &PoolManager{cipher: cipher, gatewayHost: TiDBGatewayService}
}

func (p *PoolManager) keyFromCred(cred *Credential) string {
	return strings.Join([]string{cred.Username, cred.Password, cred.Database}, "-")
}

func (p *PoolManager) GetOrCreateByCred(ctx context.Context, cred *Credential) (*DBPool, bizerr.BizError) {
	key := p.keyFromCred(cred)
	if pool, ok := p.pools.Load(key); ok && pool != nil {
		return pool.(*DBPool), nil
	}
	token, err := p.cipher.EncryptCredential(cred)
	if err != nil {
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to encrypt credential: %s", err.Error())
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	// double check to avoid race
	if pool, ok := p.pools.Load(key); ok && pool != nil {
		return pool.(*DBPool), nil
	}
	db, bizErr := p.newDatabase(ctx, cred)
	if bizErr != nil {
		return nil, bizErr
	}
	dbPool := &DBPool{db: db, cred: cred, sensitiveKey: key, expire: time.Now().Add(poolTTL), token: token}
	p.pools.Store(key, dbPool)
	return dbPool, nil
}

func (p *PoolManager) GetOrCreateByToken(ctx context.Context, token string) (*DBPool, bizerr.BizError) {
	cred, err := p.cipher.DecryptCredential(token)
	if err != nil {
		return nil, bizerr.New(bizerr.ErrCodeParams, "invalid session")
	}
	key := p.keyFromCred(cred)
	if pool, ok := p.pools.Load(key); ok && pool != nil {
		return pool.(*DBPool), nil
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	// double check to avoid race
	if pool, ok := p.pools.Load(key); ok && pool != nil {
		return pool.(*DBPool), nil
	}
	db, bizErr := p.newDatabase(ctx, cred)
	if bizErr != nil {
		return nil, bizErr
	}
	dbPool := &DBPool{db: db, cred: cred, sensitiveKey: key, expire: time.Now().Add(poolTTL), token: token}
	p.pools.Store(key, dbPool)
	return dbPool, nil
}

func (p *PoolManager) DeletePool(key string) {
	p.pools.Delete(key)
}

func (p *PoolManager) Range(fn func(string, *DBPool) bool) {
	p.pools.Range(func(key, value any) bool {
		return fn(key.(string), value.(*DBPool))
	})
}

func (p *PoolManager) newDatabase(ctx context.Context, cred *Credential) (*sql.DB, bizerr.BizError) {
	if err := ensureMysqlSetup(ctx); err != nil {
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "%s", err.Error())
	}

	cfg := mysql.NewConfig()
	cfg.DBName = cred.Database
	if cfg.DBName == "" {
		cfg.DBName = defaultDatabase
	}
	cfg.Net = "tcp"
	cfg.User = cred.Username
	cfg.Passwd = cred.Password
	cfg.Addr = net.JoinHostPort(p.gatewayHost, "4000")
	cfg.TLSConfig = "tidb"
	cfg.Params = map[string]string{"sql_select_limit": defaultLimit}
	cfg.ConnectionAttributes = fmt.Sprintf("program_name:HTTP-Gateway,%s:%s", "HTTP-Gateway-Session", fmt.Sprintf("%s/%s", cred.Username, cred.Database))

	driver, err := mysql.NewConnector(cfg)
	if err != nil {
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "%s", err.Error())
	}
	db := sql.OpenDB(driver)
	db.SetMaxOpenConns(poolMaxOpenConns)
	db.SetMaxIdleConns(poolMaxIdleConns)
	db.SetConnMaxIdleTime(poolConnMaxIdle)
	db.SetConnMaxLifetime(0)

	// https://github.com/go-sql-driver/mysql/pull/1432 will be released in v1.8.0.
	// Then we can handle the dbname with special characters.
	// databaseUrl := fmt.Sprintf("%s:%s@tcp(%s:4000)/%s?tls=tidb",
	// username, password, gatewayService, url.PathEscape(database))
	// db, err := sql.Open("mysql", databaseUrl)
	// if err != nil {
	//	return nil, err
	// }

	// ping to verify the connection is valid, otherwise throw error
	if err := db.PingContext(ctx); err != nil {
		db.Close()
		return nil, bizerr.New(bizerr.ErrCodeAuth, "%s", err.Error())
	}
	return db, nil
}

func ensureMysqlSetup(ctx context.Context) (err error) {
	mysqlSetupOnce.Do(func() {
		err = mysql.RegisterTLSConfig("tidb", &tls.Config{
			MinVersion:         tls.VersionTLS12,
			InsecureSkipVerify: true, //nolint:gosec // traffic remains inside cluster
		})
		if err != nil {
			return
		}
		err = mysql.SetLogger(&MysqlLogger{logger: log.WithContextL(ctx)})
	})
	// reset mysqlSetupOnce if error occurs to allow retry
	if err != nil {
		mysqlSetupOnce = sync.Once{}
	}
	return err
}

func (p *DBPool) Query(ctx context.Context, query string, sqlType SqlType) (*ExecuteResponse, bizerr.BizError) {
	logger := log.WithContextL(ctx).With(zap.String("token", p.token)).With(zap.String("username", p.cred.Username)).With(zap.String("database", p.cred.Database))
	logger.Info("start to query", zap.String("session type", string(StatelessSession)))
	defer func() {
		p.expire = time.Now().Add(poolTTL)
	}()

	switch sqlType {
	case Query:
		rows, err := p.db.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 := p.db.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 := p.db.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 (p *DBPool) BeginTx(txCtx context.Context, isolationLevel string) (*sql.Tx, bizerr.BizError) {
	txOptions, bizErr := getTransactionOptions(isolationLevel)
	if bizErr != nil {
		return nil, bizErr
	}
	tx, err := p.db.BeginTx(txCtx, txOptions)
	if err != nil {
		return nil, bizerr.New(ErrCodeBeginTxnFail, "%s: %s", ErrCodeBeginFailMessage, err.Error())
	}
	return tx, nil
}

func (p *DBPool) Connect(ctx context.Context) (*sql.Conn, bizerr.BizError) {
	conn, err := p.db.Conn(ctx)
	if err != nil {
		return nil, bizerr.New(ErrCodeOpenDBFail, "%s: %s", ErrCodeOpenDBFailMessage, err.Error())
	}
	return conn, nil
}