package common

import "testing"

func TestComputeRequiredDataBytes(t *testing.T) {
	const TB = uint64(1024 * 1024 * 1024 * 1024)
	const GB = uint64(1024 * 1024 * 1024)

	cases := []struct {
		name            string
		snapshot        *TiKVUsageSnapshot
		additionalNodes uint32
		threshold       uint32
		wantNonZero     bool
		verifyTarget    uint32 // expected target node count after applying required_data_bytes
	}{
		{
			name: "T4-7: 6 nodes, 2TB each, 8TB used, threshold=70, +3 nodes",
			snapshot: &TiKVUsageSnapshot{
				NodeCount:           6,
				TotalCapacityBytes:  6 * 2 * TB,
				TotalAvailableBytes: 6*2*TB - 8*TB, // 4TB available
			},
			additionalNodes: 3,
			threshold:       70,
			wantNonZero:     true,
			verifyTarget:    9,
		},
		{
			name: "3 nodes, 1TB each, 2TB used, threshold=70, +3 nodes",
			snapshot: &TiKVUsageSnapshot{
				NodeCount:           3,
				TotalCapacityBytes:  3 * TB,
				TotalAvailableBytes: 3*TB - 2*TB,
			},
			additionalNodes: 3,
			threshold:       70,
			wantNonZero:     true,
			verifyTarget:    6,
		},
		{
			name:            "nil snapshot",
			snapshot:        nil,
			additionalNodes: 3,
			threshold:       70,
			wantNonZero:     false,
		},
		{
			name: "zero node count",
			snapshot: &TiKVUsageSnapshot{
				NodeCount:           0,
				TotalCapacityBytes:  0,
				TotalAvailableBytes: 0,
			},
			additionalNodes: 3,
			threshold:       70,
			wantNonZero:     false,
		},
		{
			name: "zero threshold",
			snapshot: &TiKVUsageSnapshot{
				NodeCount:           6,
				TotalCapacityBytes:  6 * 2 * TB,
				TotalAvailableBytes: 4 * TB,
			},
			additionalNodes: 3,
			threshold:       0,
			wantNonZero:     false,
		},
		{
			name: "threshold=100 (invalid)",
			snapshot: &TiKVUsageSnapshot{
				NodeCount:           6,
				TotalCapacityBytes:  6 * 2 * TB,
				TotalAvailableBytes: 4 * TB,
			},
			additionalNodes: 3,
			threshold:       100,
			wantNonZero:     false,
		},
		{
			name: "very low usage, no additional bytes needed",
			snapshot: &TiKVUsageSnapshot{
				NodeCount:           6,
				TotalCapacityBytes:  6 * 2 * TB,
				TotalAvailableBytes: 6*2*TB - 100*GB, // only 100GB used out of 12TB
			},
			additionalNodes: 3,
			threshold:       70,
			wantNonZero:     true,
			verifyTarget:    9,
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			got := ComputeRequiredDataBytes(tc.snapshot, tc.additionalNodes, tc.threshold)

			if tc.wantNonZero && got == 0 {
				t.Errorf("expected non-zero required_data_bytes, got 0")
			}
			if !tc.wantNonZero && got != 0 {
				t.Errorf("expected 0 required_data_bytes, got %d", got)
			}

			if tc.wantNonZero && tc.verifyTarget > 0 {
				if !VerifyRequiredDataBytesYieldsTargetNodes(tc.snapshot, got, tc.threshold, tc.verifyTarget) {
					t.Errorf("computed required_data_bytes=%d does not yield target %d nodes (snapshot: nodeCount=%d, capacity=%d, available=%d)",
						got, tc.verifyTarget, tc.snapshot.NodeCount, tc.snapshot.TotalCapacityBytes, tc.snapshot.TotalAvailableBytes)
				}
			}
		})
	}
}

func TestTiKVUsageSnapshot(t *testing.T) {
	const TB = uint64(1024 * 1024 * 1024 * 1024)

	snap := &TiKVUsageSnapshot{
		NodeCount:           6,
		TotalCapacityBytes:  12 * TB,
		TotalAvailableBytes: 4 * TB,
	}

	if snap.CurrentUsageBytes() != 8*TB {
		t.Errorf("expected usage 8TB, got %d", snap.CurrentUsageBytes())
	}
	if snap.SingleNodeCapacityBytes() != 2*TB {
		t.Errorf("expected node capacity 2TB, got %d", snap.SingleNodeCapacityBytes())
	}

	pct := snap.UsagePercentage()
	if pct < 66.6 || pct > 66.7 {
		t.Errorf("expected usage ~66.67%%, got %.2f%%", pct)
	}
	if snap.RemainingBytes() != 4*TB {
		t.Errorf("expected remaining 4TB, got %d", snap.RemainingBytes())
	}
}