package orchestrator

import (
	"context"
	"fmt"
	"math/rand"
	"sync/atomic"
	"testing"
	"time"

	"cloudipvalidator/internal/db"
	"cloudipvalidator/internal/openstack"
)

// finishChecks reports a successful egress run and all three inbound sites
// for the address, so the next Tick aggregates it and releases the validator.
func finishChecks(t *testing.T, o *Orchestrator, ip *db.IPQueueItem, validatorID string) {
	t.Helper()
	ctx := context.Background()
	if err := o.RecordCheck(ctx, db.Check{
		IPID: ip.ID, IPAddress: ip.IPAddress, AttemptNumber: ip.AttemptNumber,
		ValidatorID: validatorID, Source: db.SourceEgress, CheckType: "https",
		Target: "https://example.test", Success: true, CheckedAt: db.Now(),
	}); err != nil {
		t.Fatalf("record egress check: %v", err)
	}
	if err := o.MarkEgressComplete(ctx, ip.ID); err != nil {
		t.Fatalf("mark egress complete: %v", err)
	}
	for site := 1; site <= 3; site++ {
		for _, ct := range []string{"tcp-22", "ssh", "tcp-80", "icmp"} {
			if err := o.RecordCheck(ctx, db.Check{
				IPID: ip.ID, IPAddress: ip.IPAddress, AttemptNumber: ip.AttemptNumber,
				Source: db.InboundSource(site), CheckType: ct, Target: ip.IPAddress,
				Success: true, CheckedAt: db.Now(),
			}); err != nil {
				t.Fatalf("record inbound check: %v", err)
			}
		}
		if err := o.MarkSiteComplete(ctx, ip.ID, site); err != nil {
			t.Fatalf("mark site complete: %v", err)
		}
	}
}

// The incident of 2026-10-02: a validator busy with slow checks goes silent,
// is marked unreachable, and its next heartbeat used to return it to idle
// although it still held the address. It was then handed a second address,
// whose association never ran, and both stalled until their leases expired.
func TestSilentValidatorKeepsItsAddressAfterHeartbeat(t *testing.T) {
	ctx := context.Background()
	o, d, mock := newTestOrchestrator(t, 180)
	mock.Seed("fip-a", "1.1.1.1", "svc")
	mock.Seed("fip-b", "2.2.2.2", "svc")
	if err := d.RegisterValidator(ctx, "validator-1", "host", "port-1", "v"); err != nil {
		t.Fatal(err)
	}
	if err := d.SeedQueue(ctx, []string{"1.1.1.1", "2.2.2.2"}); err != nil {
		t.Fatal(err)
	}

	o.Tick(ctx) // validator-1 claims 1.1.1.1
	a, _ := d.GetIPByAddress(ctx, "1.1.1.1")
	if err := o.SelfCheckResult(ctx, "validator-1", a.ID, true, "ok"); err != nil {
		t.Fatal(err)
	}

	// The agent goes silent for longer than heartbeat_timeout_seconds ...
	if err := d.MarkValidatorUnreachable(ctx, "validator-1"); err != nil {
		t.Fatal(err)
	}
	// ... and then speaks again while the checks are still running.
	if err := d.Heartbeat(ctx, "validator-1"); err != nil {
		t.Fatal(err)
	}
	v, _ := d.GetValidator(ctx, "validator-1")
	if v.State != db.ValidatorAssigned || v.CurrentIPID == nil || *v.CurrentIPID != a.ID {
		t.Fatalf("after heartbeat: state=%s current_ip=%v, want assigned to %d", v.State, v.CurrentIPID, a.ID)
	}

	o.Tick(ctx)
	if b, _ := d.GetIPByAddress(ctx, "2.2.2.2"); b.State != db.IPQueued {
		t.Fatalf("the busy validator was given a second address: 2.2.2.2 is %s", b.State)
	}

	// The first address finishes: only now the validator may take the next one.
	a, _ = d.GetIP(ctx, a.ID)
	finishChecks(t, o, a, "validator-1")
	o.Tick(ctx) // aggregates and releases 1.1.1.1
	o.Tick(ctx) // validator-1 is idle again and claims 2.2.2.2
	if b, _ := d.GetIPByAddress(ctx, "2.2.2.2"); b.State != db.IPAwaitingSelfCheck {
		t.Fatalf("2.2.2.2 is %s, want awaiting_self_check once the validator is free", b.State)
	}
}

// A dead validator whose lease was reclaimed must stay out of rotation until
// it speaks again; before, reclaiming set it idle and it was handed a fresh
// address every lease period, burning each address's retries.
func TestUnreachableValidatorGetsNoAddressesAfterLeaseReclaim(t *testing.T) {
	ctx := context.Background()
	o, d, mock := newTestOrchestrator(t, 1)
	mock.Seed("fip-a", "1.1.1.1", "svc")
	mock.Seed("fip-b", "2.2.2.2", "svc")
	_ = d.RegisterValidator(ctx, "validator-1", "host", "port-1", "v")
	_ = d.SeedQueue(ctx, []string{"1.1.1.1", "2.2.2.2"})

	o.Tick(ctx) // claims 1.1.1.1, the agent never answers
	if err := d.MarkValidatorUnreachable(ctx, "validator-1"); err != nil {
		t.Fatal(err)
	}
	time.Sleep(1100 * time.Millisecond)
	o.Cfg.LeaseTTLSeconds = 180
	o.Tick(ctx) // lease sweep reclaims 1.1.1.1

	a, _ := d.GetIPByAddress(ctx, "1.1.1.1")
	if a.RetryCount != 1 {
		t.Fatalf("retry_count = %d, want the lease to be reclaimed once", a.RetryCount)
	}
	v, _ := d.GetValidator(ctx, "validator-1")
	if v.State != db.ValidatorUnreachable || v.CurrentIPID != nil {
		t.Fatalf("validator state=%s current_ip=%v, want unreachable and empty", v.State, v.CurrentIPID)
	}
	o.Tick(ctx)
	if a, _ := d.GetIPByAddress(ctx, "1.1.1.1"); a.State != db.IPQueued {
		t.Fatalf("an unreachable validator was handed an address: 1.1.1.1 is %s", a.State)
	}

	if err := d.Heartbeat(ctx, "validator-1"); err != nil { // it is back
		t.Fatal(err)
	}
	if v, _ := d.GetValidator(ctx, "validator-1"); v.State != db.ValidatorIdle {
		t.Fatalf("after heartbeat state=%s, want idle (it holds nothing)", v.State)
	}
	o.Tick(ctx)
	if a, _ := d.GetIPByAddress(ctx, "1.1.1.1"); a.State != db.IPAwaitingSelfCheck {
		t.Fatalf("1.1.1.1 is %s, want it picked up again", a.State)
	}
}

// countingOS counts Disassociate calls.
type countingOS struct {
	*openstack.MockClient
	disassociations int32
}

func (c *countingOS) DisassociateFloatingIP(ctx context.Context, fipID string) error {
	atomic.AddInt32(&c.disassociations, 1)
	return c.MockClient.DisassociateFloatingIP(ctx, fipID)
}

// Finished addresses keep their fip_id for display but their floating IP is
// already free; clearing the queue must not make a cloud call for each of
// them (on 2026-10-02 that was >1000 sequential calls: 256 s).
func TestClearQueueDoesNotTouchFinishedAddresses(t *testing.T) {
	ctx := context.Background()
	o, d, mock := newTestOrchestrator(t, 180)
	cos := &countingOS{MockClient: mock}
	o.OS = cos

	const finished = 300
	var addrs []string
	for i := 0; i < finished; i++ {
		addrs = append(addrs, fmt.Sprintf("10.9.%d.%d", i/200, i%200+1))
	}
	if err := d.SeedQueue(ctx, addrs); err != nil {
		t.Fatal(err)
	}
	for i, a := range addrs {
		state := []string{db.IPDone, db.IPFailed, db.IPOccupied}[i%3]
		if _, err := d.ExecContext(ctx, `UPDATE ip_queue SET state=?, fip_id=? WHERE ip_address=?`, state, fmt.Sprintf("fip-old-%d", i), a); err != nil {
			t.Fatal(err)
		}
	}
	// One address really holds a floating IP.
	mock.Seed("fip-live", "1.2.3.4", "svc")
	_ = d.RegisterValidator(ctx, "validator-1", "host", "port-1", "v")
	_ = d.SeedQueue(ctx, []string{"1.2.3.4"})
	live, _ := d.GetIPByAddress(ctx, "1.2.3.4")
	if _, err := d.ExecContext(ctx, `UPDATE ip_queue SET sequence=-1 WHERE id=?`, live.ID); err != nil {
		t.Fatal(err)
	}
	o.Tick(ctx) // validator-1 takes the live address (lowest sequence) and attaches it
	if got := portFIPs(t, mock, "port-1"); len(got) != 1 {
		t.Fatalf("setup: the live address is not attached (%d floating ips on port-1)", len(got))
	}
	atomic.StoreInt32(&cos.disassociations, 0)

	start := time.Now()
	if _, err := o.ClearQueue(ctx); err != nil {
		t.Fatalf("clear queue: %v", err)
	}
	// One call for the live address; the port sweep finds nothing left.
	if n := atomic.LoadInt32(&cos.disassociations); n != 1 {
		t.Fatalf("clear queue made %d disassociate calls, want 1 (finished addresses must be skipped)", n)
	}
	if got := portFIPs(t, mock, "port-1"); len(got) != 0 {
		t.Fatalf("the live floating ip is still attached: %+v", got)
	}
	if left, _ := d.ListIPs(ctx); len(left) != 0 {
		t.Fatalf("%d rows left after clear", len(left))
	}
	t.Logf("clear of %d rows took %s", finished+1, time.Since(start))
}

// A bulk detach runs in parallel but never above maxParallelDetach at once.
func TestDisassociateAllIsBoundedAndParallel(t *testing.T) {
	o, _, mock := newTestOrchestrator(t, 180)
	g := &gaugeOS{MockClient: mock, delay: 30 * time.Millisecond}
	o.OS = g
	var refs []db.FIPRef
	for i := 0; i < 40; i++ {
		id := fmt.Sprintf("fip-%d", i)
		mock.Seed(id, fmt.Sprintf("10.8.0.%d", i+1), "svc")
		refs = append(refs, db.FIPRef{IPID: int64(i), IPAddress: id, FIPID: id})
	}
	start := time.Now()
	o.disassociateAll(context.Background(), refs, "test")
	elapsed := time.Since(start)
	if peak := atomic.LoadInt32(&g.peak); peak < 2 || peak > maxParallelDetach {
		t.Fatalf("peak concurrency %d, want between 2 and %d", peak, maxParallelDetach)
	}
	if elapsed > 600*time.Millisecond { // 40 x 30 ms sequentially is 1.2 s
		t.Fatalf("took %s: the detach is not parallel", elapsed)
	}
}

type gaugeOS struct {
	*openstack.MockClient
	delay     time.Duration
	cur, peak int32
}

func (g *gaugeOS) DisassociateFloatingIP(ctx context.Context, fipID string) error {
	n := atomic.AddInt32(&g.cur, 1)
	for {
		p := atomic.LoadInt32(&g.peak)
		if n <= p || atomic.CompareAndSwapInt32(&g.peak, p, n) {
			break
		}
	}
	time.Sleep(g.delay)
	atomic.AddInt32(&g.cur, -1)
	return g.MockClient.DisassociateFloatingIP(ctx, fipID)
}

// Rows that disagree with the queue are repaired by ReconcileValidators (run on every Tick).
func TestReconcileRepairsInconsistentValidators(t *testing.T) {
	ctx := context.Background()
	_, d, _ := newTestOrchestrator(t, 180)
	_ = d.RegisterValidator(ctx, "validator-1", "host", "port-1", "v")
	_ = d.RegisterValidator(ctx, "validator-2", "host", "port-2", "v")
	_ = d.SeedQueue(ctx, []string{"1.1.1.1", "2.2.2.2"})

	finished, _ := d.GetIPByAddress(ctx, "1.1.1.1")
	foreign, _ := d.GetIPByAddress(ctx, "2.2.2.2")
	// validator-1 still points at a finished address; validator-2 at an
	// address owned by somebody else (what an older version could leave behind).
	for _, q := range []string{
		fmt.Sprintf(`UPDATE ip_queue SET state='done' WHERE id=%d`, finished.ID),
		fmt.Sprintf(`UPDATE ip_queue SET state='checking', owner_validator_id='validator-1' WHERE id=%d`, foreign.ID),
		fmt.Sprintf(`UPDATE validators SET state='assigned', current_ip_id=%d WHERE validator_id='validator-1'`, finished.ID),
		fmt.Sprintf(`UPDATE validators SET state='assigned', current_ip_id=%d WHERE validator_id='validator-2'`, foreign.ID),
	} {
		if _, err := d.ExecContext(ctx, q); err != nil {
			t.Fatal(err)
		}
	}
	n, err := d.ReconcileValidators(ctx)
	if err != nil || n != 2 {
		t.Fatalf("reconcile repaired %d validators (err %v), want 2", n, err)
	}
	for _, id := range []string{"validator-1", "validator-2"} {
		v, _ := d.GetValidator(ctx, id)
		if v.State != db.ValidatorIdle || v.CurrentIPID != nil {
			t.Fatalf("%s: state=%s current_ip=%v, want idle", id, v.State, v.CurrentIPID)
		}
	}
	// A consistent validator is left alone.
	if n, _ := d.ReconcileValidators(ctx); n != 0 {
		t.Fatalf("second reconcile repaired %d, want 0", n)
	}
}

// Randomised run: many validators, random silences and association failures.
// After every tick each validator holds at most one address, the validator
// and the address agree on who holds what, and no lease is ever reclaimed.
func TestRandomFlowKeepsOneAddressPerValidator(t *testing.T) {
	ctx := context.Background()
	rng := rand.New(rand.NewSource(42))
	o, d, mock := newTestOrchestrator(t, 180)

	const validators, addresses = 12, 60
	var addrs []string
	for i := 0; i < validators; i++ {
		_ = d.RegisterValidator(ctx, fmt.Sprintf("validator-%d", i+1), "host", fmt.Sprintf("port-%d", i+1), "v")
	}
	for i := 0; i < addresses; i++ {
		a := fmt.Sprintf("10.7.%d.%d", i/200, i%200+1)
		addrs = append(addrs, a)
		mock.Seed(fmt.Sprintf("fip-%d", i), a, "svc")
	}
	if err := d.SeedQueue(ctx, addrs); err != nil {
		t.Fatal(err)
	}

	checkInvariants := func(tick int) {
		t.Helper()
		vs, _ := d.ListValidators(ctx)
		holders := map[int64]string{}
		for _, v := range vs {
			if v.CurrentIPID == nil {
				if v.State == db.ValidatorAssigned {
					t.Fatalf("tick %d: %s is assigned but holds nothing", tick, v.ValidatorID)
				}
				continue
			}
			ip, err := d.GetIP(ctx, *v.CurrentIPID)
			if err != nil {
				t.Fatalf("tick %d: %s points at a missing address %d", tick, v.ValidatorID, *v.CurrentIPID)
			}
			if ip.OwnerValidatorID == nil || *ip.OwnerValidatorID != v.ValidatorID {
				t.Fatalf("tick %d: %s holds %s but its owner is %v", tick, v.ValidatorID, ip.IPAddress, ip.OwnerValidatorID)
			}
			if ip.State == db.IPDone || ip.State == db.IPFailed || ip.State == db.IPOccupied {
				t.Fatalf("tick %d: %s still holds the finished address %s", tick, v.ValidatorID, ip.IPAddress)
			}
			if prev, ok := holders[ip.ID]; ok {
				t.Fatalf("tick %d: %s is held by both %s and %s", tick, ip.IPAddress, prev, v.ValidatorID)
			}
			holders[ip.ID] = v.ValidatorID
		}
		// Every owned, unfinished address is the current one of its owner.
		ips, _ := d.ListIPs(ctx)
		perOwner := map[string]int{}
		for _, ip := range ips {
			if ip.OwnerValidatorID != nil && ip.State != db.IPDone && ip.State != db.IPFailed && ip.State != db.IPOccupied {
				perOwner[*ip.OwnerValidatorID]++
			}
		}
		for owner, n := range perOwner {
			if n > 1 {
				t.Fatalf("tick %d: %s owns %d unfinished addresses", tick, owner, n)
			}
		}
	}

	checkTicks := map[int64]int{} // address id -> ticks spent in checking
	done := false
	for tick := 1; tick <= 600 && !done; tick++ {
		// Occasionally an association fails (409 on a port).
		if rng.Intn(15) == 0 {
			mock.AssociateFailures = map[string]error{fmt.Sprintf("fip-%d", rng.Intn(addresses)): fmt.Errorf("409 conflict")}
		}
		o.Tick(ctx)

		vs, _ := d.ListValidators(ctx)
		for _, v := range vs {
			// A busy agent sometimes goes silent, then speaks again.
			if v.CurrentIPID != nil && v.State == db.ValidatorAssigned && rng.Intn(10) == 0 {
				_ = d.MarkValidatorUnreachable(ctx, v.ValidatorID)
			}
			if v.State == db.ValidatorUnreachable && rng.Intn(3) == 0 {
				_ = d.Heartbeat(ctx, v.ValidatorID)
			}
			item, _, err := o.AssignmentForValidator(ctx, v.ValidatorID)
			if err != nil || item == nil {
				continue
			}
			if item.State == db.IPAwaitingSelfCheck {
				_ = o.SelfCheckResult(ctx, v.ValidatorID, item.ID, true, "ok")
				continue
			}
			checkTicks[item.ID]++
			if checkTicks[item.ID] >= 1+rng.Intn(4) {
				finishChecks(t, o, item, v.ValidatorID)
				delete(checkTicks, item.ID)
			}
		}
		checkInvariants(tick)

		counts, _, _ := d.CountIPsByState(ctx)
		unfinished := 0
		for st, n := range counts {
			if st != db.IPDone && st != db.IPFailed && st != db.IPOccupied {
				unfinished += n
			}
		}
		done = unfinished == 0
	}
	if !done {
		counts, _, _ := d.CountIPsByState(ctx)
		t.Fatalf("not all addresses finished: %v", counts)
	}
	var expired int
	if err := d.QueryRowContext(ctx, `SELECT COUNT(*) FROM events WHERE event_type='lease_expired'`).Scan(&expired); err != nil {
		t.Fatal(err)
	}
	if expired != 0 {
		t.Fatalf("%d leases were reclaimed, want 0", expired)
	}
}