package kvm2

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"

	"gitea.home.hrajfrisbee.cz/kacerr/egress-proxies-operator/internal/provider"
)

var (
	hostA = provider.KVMHost{Name: "beelink", Address: "beelink:22", User: "virt"}
	hostB = provider.KVMHost{Name: "lab2", Address: "lab2:22", User: "virt"}
)

func TestHosts_oneDialPerHostUnderConcurrency(t *testing.T) {
	t.Parallel()
	shared := newFakeBackend()
	d := &stubDialer{backend: shared}
	table := newHostTable(d)

	const n = 50
	var wg sync.WaitGroup
	got := make([]backend, n)
	errs := make([]error, n)
	for i := range n {
		wg.Add(1)
		go func() {
			defer wg.Done()
			got[i], errs[i] = table.get(context.Background(), hostA)
		}()
	}
	wg.Wait()
	for i := range n {
		if errs[i] != nil {
			t.Fatalf("get #%d error = %v", i, errs[i])
		}
		if got[i] != backend(shared) {
			t.Fatalf("get #%d returned a different backend", i)
		}
	}
	if d.dials.Load() != 1 {
		t.Errorf("dials = %d, want 1", d.dials.Load())
	}
}

func TestHosts_redialsAfterDrop(t *testing.T) {
	t.Parallel()
	first := newFakeBackend()
	second := newFakeBackend()
	backends := []backend{first, second}
	d := &stubDialer{fresh: func() backend {
		b := backends[0]
		backends = backends[1:]
		return b
	}}
	table := newHostTable(d)

	b, err := table.get(context.Background(), hostA)
	if err != nil || b != backend(first) {
		t.Fatalf("first get = %v, %v, want the first backend", b, err)
	}
	if b, _ := table.get(context.Background(), hostA); b != backend(first) {
		t.Fatal("second get did not reuse the live backend")
	}
	first.markDead()
	b, err = table.get(context.Background(), hostA)
	if err != nil || b != backend(second) {
		t.Fatalf("get after drop = %v, %v, want the second backend", b, err)
	}
	if first.closed != 1 {
		t.Errorf("dropped backend closed %d times, want 1", first.closed)
	}
	if d.dials.Load() != 2 {
		t.Errorf("dials = %d, want 2", d.dials.Load())
	}
}

func TestHosts_dialFreshIsNotCached(t *testing.T) {
	t.Parallel()
	d := &stubDialer{fresh: func() backend { return newFakeBackend() }}
	table := newHostTable(d)

	pooled, err := table.get(context.Background(), hostA)
	if err != nil {
		t.Fatalf("get error = %v", err)
	}
	fresh, err := table.dialFresh(context.Background(), hostA)
	if err != nil {
		t.Fatalf("dialFresh error = %v", err)
	}
	if fresh == pooled {
		t.Fatal("dialFresh returned the pooled backend")
	}
	again, _ := table.get(context.Background(), hostA)
	if again != pooled {
		t.Error("get after dialFresh did not return the pooled backend")
	}
	if d.dials.Load() != 2 {
		t.Errorf("dials = %d, want 2", d.dials.Load())
	}
	if fresh.(*fakeBackend).closed != 0 {
		t.Error("dialFresh's backend was closed by the table; the caller owns it")
	}
}

func TestHosts_waitingForDialEndsWithContext(t *testing.T) {
	t.Parallel()
	release := make(chan struct{})
	d := &stubDialer{block: release}
	table := newHostTable(d)

	firstDone := make(chan error, 1)
	go func() {
		_, err := table.get(context.Background(), hostA)
		firstDone <- err
	}()
	waitForDials(t, d, 1)

	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	start := time.Now()
	_, err := table.get(ctx, hostA)
	if time.Since(start) > time.Second {
		t.Errorf("second get took %v, want to give up with its ctx", time.Since(start))
	}
	if !errors.Is(err, provider.ErrTransient) || !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("second get error = %v, want ErrTransient wrapping DeadlineExceeded", err)
	}
	close(release)
	if err := <-firstDone; err != nil {
		t.Errorf("first get error = %v, want nil once the dial completes", err)
	}
	if d.dials.Load() != 1 {
		t.Errorf("dials = %d, want 1: the waiter must not dial itself", d.dials.Load())
	}
}

func TestHosts_hostsDialIndependently(t *testing.T) {
	t.Parallel()
	release := make(chan struct{})
	defer close(release)
	dA := &stubDialer{block: release}
	dB := &stubDialer{}
	table := newHostTable(perHostDialer{hostA.Name: dA, hostB.Name: dB})

	go func() { _, _ = table.get(context.Background(), hostA) }()
	waitForDials(t, dA, 1)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	start := time.Now()
	if _, err := table.get(ctx, hostB); err != nil {
		t.Fatalf("get(hostB) error = %v while hostA's dial is blocked", err)
	}
	if time.Since(start) > 500*time.Millisecond {
		t.Errorf("get(hostB) took %v, want not to wait for hostA", time.Since(start))
	}
}

func TestHosts_dialErrorIsNotCached(t *testing.T) {
	t.Parallel()
	d := &stubDialer{err: transientf("boom")}
	table := newHostTable(d)
	if _, err := table.get(context.Background(), hostA); !errors.Is(err, provider.ErrTransient) {
		t.Fatalf("get error = %v, want the dial error", err)
	}
	d.err = nil
	if _, err := table.get(context.Background(), hostA); err != nil {
		t.Fatalf("get after the dialer recovered error = %v", err)
	}
	if d.dials.Load() != 2 {
		t.Errorf("dials = %d, want 2", d.dials.Load())
	}
}

// perHostDialer routes each host to its own stub.
type perHostDialer map[string]*stubDialer

func (p perHostDialer) dial(ctx context.Context, host provider.KVMHost) (backend, error) {
	return p[host.Name].dial(ctx, host)
}

func waitForDials(t *testing.T, d *stubDialer, n int32) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for d.dials.Load() < n {
		if time.Now().After(deadline) {
			t.Fatalf("dials = %d, want %d", d.dials.Load(), n)
		}
		time.Sleep(time.Millisecond)
	}
}