package executor

import (
	"context"
	"database/sql"
	"errors"
	"strings"
	"testing"
	"time"

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

type countingLeaseProvider struct {
	calls int
	kinds []mysqlbackend.ConnLeaseKind
}

func (p *countingLeaseProvider) AcquireConnLease(kind mysqlbackend.ConnLeaseKind) (func(), error) {
	p.calls++
	p.kinds = append(p.kinds, kind)
	return func() {}, nil
}

func TestEnsureConnForInferenceUsesExistingConn(t *testing.T) {
	exec := &Executor{}
	dummyConn := &sql.Conn{}
	ctx := mysqlbackend.WithConn(context.Background(), dummyConn)

	newCtx, release, err := exec.ensureConnForInference(ctx, "user", "pass")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if release != nil {
		t.Fatalf("expected nil release when connection already present")
	}
	conn, ok := mysqlbackend.ConnFromContext(newCtx)
	if !ok || conn != dummyConn {
		t.Fatalf("expected existing conn to remain in context")
	}
}

func TestEnsureConnForInferenceWithoutPoolFails(t *testing.T) {
	exec := &Executor{}
	_, _, err := exec.ensureConnForInference(context.Background(), "user", "pass")
	if err == nil {
		t.Fatalf("expected error without pool manager")
	}
}

func TestGetConnSkipsLeaseWhenHeldStandardLease(t *testing.T) {
	exec := &Executor{}
	provider := &countingLeaseProvider{}
	ctx := mysqlbackend.WithConnLeaseProvider(context.Background(), provider)
	ctx = mysqlbackend.WithHeldConnLease(ctx, mysqlbackend.ConnLeaseKindStandard)

	_, _, err := exec.getConn(ctx, "user", "pass")
	if err == nil || !strings.Contains(err.Error(), "pool manager is not configured") {
		t.Fatalf("expected pool manager error, got %v", err)
	}
	if provider.calls != 0 {
		t.Fatalf("expected held standard lease to skip provider acquisition, got %d calls", provider.calls)
	}
}

func TestGetConnUsesPreferredBulkLeaseKind(t *testing.T) {
	exec := &Executor{}
	provider := &countingLeaseProvider{}
	ctx := mysqlbackend.WithConnLeaseProvider(context.Background(), provider)
	ctx = mysqlbackend.WithPreferredConnLeaseKind(ctx, mysqlbackend.ConnLeaseKindBulk)

	_, _, err := exec.getConn(ctx, "user", "pass")
	if err == nil || !strings.Contains(err.Error(), "pool manager is not configured") {
		t.Fatalf("expected pool manager error, got %v", err)
	}
	if provider.calls != 1 {
		t.Fatalf("expected one provider call, got %d", provider.calls)
	}
	if len(provider.kinds) != 1 || provider.kinds[0] != mysqlbackend.ConnLeaseKindBulk {
		t.Fatalf("expected preferred bulk lease kind, got %v", provider.kinds)
	}
}

func TestAcquireUnloggedBatchConnLeaseEnforcesLimit(t *testing.T) {
	exec := &Executor{unloggedBatchConnLimiter: newConnLimiter(1)}

	release, err := exec.acquireUnloggedBatchConnLease(context.Background())
	if err != nil {
		t.Fatalf("first bulk lease failed: %v", err)
	}
	if release == nil {
		t.Fatal("expected release func for bulk lease")
	}

	waitDone := make(chan struct{})
	go func() {
		secondRelease, secondErr := exec.acquireUnloggedBatchConnLease(context.Background())
		if secondErr != nil {
			t.Errorf("second bulk lease failed: %v", secondErr)
			close(waitDone)
			return
		}
		if secondRelease != nil {
			secondRelease()
		}
		close(waitDone)
	}()

	select {
	case <-waitDone:
		t.Fatal("expected second bulk lease to wait for release")
	case <-time.After(50 * time.Millisecond):
	}

	release()

	select {
	case <-waitDone:
	case <-time.After(200 * time.Millisecond):
		t.Fatal("expected waiting bulk lease to proceed after release")
	}

	heldRelease, err := exec.acquireUnloggedBatchConnLease(context.Background())
	if err != nil {
		t.Fatalf("reacquire bulk lease failed: %v", err)
	}
	canceledCtx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := exec.acquireUnloggedBatchConnLease(canceledCtx); !errors.Is(err, context.Canceled) {
		if heldRelease != nil {
			heldRelease()
		}
		t.Fatalf("expected context cancellation, got %v", err)
	}
	if heldRelease != nil {
		heldRelease()
	}
}