package rules

import (
	"os"
	"path/filepath"
	"testing"
	"time"
)

func TestConnectionStoreConnectAndRecommend(t *testing.T) {
	tmpDir := t.TempDir()
	store := NewConnectionStore(tmpDir)
	if err := store.Load(); err != nil {
		t.Fatal(err)
	}

	availableIPs := []string{"5.180.97.198", "5.180.97.199", "5.180.97.197"}
	healthyIPs := map[string]bool{"5.180.97.198": true, "5.180.97.199": true, "5.180.97.197": true}

	// Studio 1 connects to 198
	store.Connect("1.2.3.4", "5.180.97.198", "nl-198", "windows", "2.0.11")

	// Studio 1 asks — should get 198 (least loaded: 1 client vs 0/0)
	// Actually: 198=1, 199=0, 197=0 → should get 199 or 197 (least loaded)
	rec1 := store.GetRecommendation("1.2.3.4", availableIPs, healthyIPs)
	if rec1.RecommendedServerIP == "5.180.97.198" {
		// This is OK if load balancing picks a different server
		t.Logf("studio 1 recommended: %s (reason: %s)", rec1.RecommendedServerIP, rec1.Reason)
	} else {
		t.Logf("studio 1 recommended different server: %s (load-based)", rec1.RecommendedServerIP)
	}

	// Studio 2 is new — should also get least loaded
	rec2 := store.GetRecommendation("9.9.9.9", availableIPs, healthyIPs)
	if rec2.RecommendedServerIP == "" {
		t.Fatal("expected recommendation for new studio")
	}
	t.Logf("studio 2 recommended: %s (reason: %s)", rec2.RecommendedServerIP, rec2.Reason)
}

func TestPureLoadBalancing(t *testing.T) {
	tmpDir := t.TempDir()
	store := NewConnectionStore(tmpDir)
	if err := store.Load(); err != nil {
		t.Fatal(err)
	}

	availableIPs := []string{"5.180.97.198", "5.180.97.199", "5.180.97.197"}
	healthyIPs := map[string]bool{"5.180.97.198": true, "5.180.97.199": true, "5.180.97.197": true}

	// 3 studios connect to 198 (overload it)
	for i := 0; i < 3; i++ {
		ip := "10.0.0." + string(rune('1'+i))
		store.Connect(ip, "5.180.97.198", "nl-198", "windows", "")
	}

	// New studio should NOT get 198 (3 clients) — should get 199 or 197 (0 clients)
	rec := store.GetRecommendation("99.99.99.99", availableIPs, healthyIPs)
	if rec.RecommendedServerIP == "5.180.97.198" {
		t.Fatalf("should not recommend overloaded server, got %s", rec.RecommendedServerIP)
	}
	t.Logf("new studio recommended: %s (reason: %s)", rec.RecommendedServerIP, rec.Reason)

	// Even studio 1 (home=198) should get load-balanced recommendation
	rec2 := store.GetRecommendation("10.0.0.1", availableIPs, healthyIPs)
	t.Logf("studio 1 re-recommended: %s (reason: %s, isRebalance: %v)",
		rec2.RecommendedServerIP, rec2.Reason, rec2.IsRebalance)
}

func TestHomeServerUnhealthy(t *testing.T) {
	tmpDir := t.TempDir()
	store := NewConnectionStore(tmpDir)
	if err := store.Load(); err != nil {
		t.Fatal(err)
	}

	availableIPs := []string{"5.180.97.198", "5.180.97.199"}
	// 198 is NOT healthy
	healthyIPs := map[string]bool{"5.180.97.199": true}

	// Studio 1 has connected to 198
	store.Connect("1.2.3.4", "5.180.97.198", "nl-198", "windows", "")

	// But 198 is unhealthy — should recommend 199
	rec := store.GetRecommendation("1.2.3.4", availableIPs, healthyIPs)
	if rec.RecommendedServerIP == "5.180.97.198" {
		t.Fatalf("should not recommend unhealthy server, got %s", rec.RecommendedServerIP)
	}
	if rec.RecommendedServerIP != "5.180.97.199" {
		t.Fatalf("should recommend healthy server 199, got %s", rec.RecommendedServerIP)
	}
}

func TestDisconnect(t *testing.T) {
	tmpDir := t.TempDir()
	store := NewConnectionStore(tmpDir)
	if err := store.Load(); err != nil {
		t.Fatal(err)
	}

	availableIPs := []string{"5.180.97.198"}

	store.Connect("1.2.3.4", "5.180.97.198", "nl-198", "windows", "")

	load := store.GetLoadInfo(availableIPs)
	if len(load) == 0 || load[0].ActiveClients != 1 {
		t.Fatalf("expected 1 active client, got %v", load)
	}

	store.Disconnect("1.2.3.4")

	load = store.GetLoadInfo(availableIPs)
	if len(load) == 0 || load[0].ActiveClients != 0 {
		t.Fatalf("expected 0 active clients after disconnect, got %v", load)
	}
}

func TestSessionExpiry(t *testing.T) {
	tmpDir := t.TempDir()
	store := NewConnectionStore(tmpDir)
	store.staleAfter = 1 * time.Millisecond
	if err := store.Load(); err != nil {
		t.Fatal(err)
	}

	availableIPs := []string{"5.180.97.198"}
	healthyIPs := map[string]bool{"5.180.97.198": true}

	store.Connect("1.2.3.4", "5.180.97.198", "nl-198", "windows", "")
	time.Sleep(10 * time.Millisecond)

	rec := store.GetRecommendation("1.2.3.4", availableIPs, healthyIPs)
	if rec.RecommendedServerIP != "5.180.97.198" {
		t.Fatalf("expected recommendation to 198 after session expiry, got %s", rec.RecommendedServerIP)
	}

	load := store.GetLoadInfo(availableIPs)
	if len(load) == 0 || load[0].ActiveClients != 0 {
		t.Fatalf("expected 0 active clients after expiry, got %v", load)
	}
}

func TestPersistence(t *testing.T) {
	tmpDir := t.TempDir()

	store1 := NewConnectionStore(tmpDir)
	store1.Connect("1.2.3.4", "5.180.97.199", "nl-199", "windows", "")

	store2 := NewConnectionStore(tmpDir)
	if err := store2.Load(); err != nil {
		t.Fatal(err)
	}

	availableIPs := []string{"5.180.97.199"}
	healthyIPs := map[string]bool{"5.180.97.199": true}
	rec := store2.GetRecommendation("1.2.3.4", availableIPs, healthyIPs)
	if rec.RecommendedServerIP != "5.180.97.199" {
		t.Fatalf("expected recommendation to 199, got %s", rec.RecommendedServerIP)
	}

	_, err := os.Stat(filepath.Join(tmpDir, "connections.json"))
	if err != nil {
		t.Fatal("expected connections.json to exist")
	}
}

func TestLoadInfoFormat(t *testing.T) {
	tmpDir := t.TempDir()
	store := NewConnectionStore(tmpDir)
	store.maxCap = 10

	store.Connect("1.1.1.1", "5.180.97.198", "nl-198", "windows", "")
	store.Connect("2.2.2.2", "5.180.97.198", "nl-198", "windows", "")
	store.Connect("3.3.3.3", "5.180.97.199", "nl-199", "linux", "")

	availableIPs := []string{"5.180.97.198", "5.180.97.199"}
	load := store.GetLoadInfo(availableIPs)

	if len(load) != 2 {
		t.Fatalf("expected 2 server load entries, got %d", len(load))
	}

	for _, info := range load {
		if info.ServerIP == "5.180.97.198" {
			if info.ActiveClients != 2 {
				t.Errorf("expected 2 clients on 198, got %d", info.ActiveClients)
			}
			if info.LoadPercent != 20 {
				t.Errorf("expected 20%% load on 198, got %d", info.LoadPercent)
			}
		}
		if info.ServerIP == "5.180.97.199" {
			if info.ActiveClients != 1 {
				t.Errorf("expected 1 client on 199, got %d", info.ActiveClients)
			}
			if info.LoadPercent != 10 {
				t.Errorf("expected 10%% load on 199, got %d", info.LoadPercent)
			}
		}
	}
}