package executor

import (
	"context"
	"database/sql"
	"fmt"
	"sync"

	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/cqlerrors"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/mysqlbackend"
)

type connLimiter struct {
	ch chan struct{}
}

func newConnLimiter(max int) *connLimiter {
	if max <= 0 {
		return nil
	}
	return &connLimiter{ch: make(chan struct{}, max)}
}

func (l *connLimiter) Acquire(ctx context.Context) error {
	if l == nil {
		return nil
	}
	select {
	case l.ch <- struct{}{}:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (l *connLimiter) Release() {
	if l == nil {
		return
	}
	select {
	case <-l.ch:
	default:
	}
}

func joinConnReleases(conn *sql.Conn, releases ...func()) func() {
	var active []func()
	for _, release := range releases {
		if release != nil {
			active = append(active, release)
		}
	}
	if conn == nil && len(active) == 0 {
		return nil
	}
	var once sync.Once
	return func() {
		once.Do(func() {
			if conn != nil {
				_ = conn.Close()
			}
			for _, release := range active {
				release()
			}
		})
	}
}

func (e *Executor) acquireUnloggedBatchConnLease(ctx context.Context) (func(), error) {
	if e == nil || e.unloggedBatchConnLimiter == nil {
		return nil, nil
	}
	if ctx == nil {
		ctx = context.Background()
	}
	if err := e.unloggedBatchConnLimiter.Acquire(ctx); err != nil {
		if err == context.Canceled || err == context.DeadlineExceeded {
			return nil, err
		}
		return nil, cqlerrors.Overloaded("unlogged batch MySQL connection limit exceeded")
	}
	var once sync.Once
	return func() {
		once.Do(func() {
			e.unloggedBatchConnLimiter.Release()
		})
	}, nil
}

func (e *Executor) getConn(ctx context.Context, user, password string) (*sql.Conn, func(), error) {
	if conn, ok := mysqlbackend.ConnFromContext(ctx); ok {
		return conn, nil, nil
	}
	leaseRelease, err := mysqlbackend.AcquireConnLeaseFromContext(ctx, mysqlbackend.ConnLeaseKindStandard)
	if err != nil {
		return nil, nil, err
	}
	if e.poolManager == nil {
		if leaseRelease != nil {
			leaseRelease()
		}
		return nil, nil, fmt.Errorf("pool manager is not configured")
	}
	pool, err := e.poolManager.GetPool(user, password)
	if err != nil {
		if leaseRelease != nil {
			leaseRelease()
		}
		return nil, nil, fmt.Errorf("failed to get connection pool: %w", err)
	}
	conn, err := pool.Conn(ctx)
	if err != nil {
		if leaseRelease != nil {
			leaseRelease()
		}
		return nil, nil, fmt.Errorf("failed to get connection: %w", err)
	}
	return conn, joinConnReleases(conn, leaseRelease), nil
}

// ValidateCredentials verifies that the user/password can open a TiDB connection.
func (e *Executor) ValidateCredentials(ctx context.Context, user, password string) error {
	if e == nil || e.poolManager == nil {
		return fmt.Errorf("pool manager is not configured")
	}
	return e.poolManager.ValidateCredentials(ctx, user, password)
}

func (e *Executor) ensureConnForInference(ctx context.Context, user, password string) (context.Context, func(), error) {
	if ctx == nil {
		ctx = context.Background()
	}
	conn, release, err := e.getConn(ctx, user, password)
	if err != nil {
		return ctx, nil, err
	}
	if release == nil || conn == nil {
		return ctx, nil, nil
	}
	ctx = mysqlbackend.WithConn(ctx, conn)
	return ctx, release, nil
}