package storage

import (
	"testing"
)

func TestNewStationNeighborsTable(t *testing.T) {
	snt := NewStationNeighborsTable()
	if snt.neighbors == nil {
		t.Error("expected neighbors map to be initialized")
	}
}

func TestStationNeighborsTableAdd(t *testing.T) {
	snt := NewStationNeighborsTable()
	snt.Add("c1", "s1", "Station One", "geo")
	snt.Add("c1", "s2", "Station Two", "manual")

	neighbors := snt.GetByCity("c1")
	if len(neighbors) != 2 {
		t.Errorf("expected 2 neighbors, got %d", len(neighbors))
	}

	if neighbors[0].StationID != "s1" {
		t.Errorf("expected first neighbor StationID 's1', got '%s'", neighbors[0].StationID)
	}
	if neighbors[0].Source != "geo" {
		t.Errorf("expected first neighbor Source 'geo', got '%s'", neighbors[0].Source)
	}

	if neighbors[1].StationID != "s2" {
		t.Errorf("expected second neighbor StationID 's2', got '%s'", neighbors[1].StationID)
	}
	if neighbors[1].Source != "manual" {
		t.Errorf("expected second neighbor Source 'manual', got '%s'", neighbors[1].Source)
	}
}

func TestStationNeighborsTableGetByCity(t *testing.T) {
	snt := NewStationNeighborsTable()
	snt.Add("c1", "s1", "Station One", "geo")

	neighbors := snt.GetByCity("c1")
	if len(neighbors) != 1 {
		t.Errorf("expected 1 neighbor for c1, got %d", len(neighbors))
	}

	neighborsEmpty := snt.GetByCity("c999")
	if neighborsEmpty != nil {
		t.Errorf("expected nil for non-existent city, got %v", neighborsEmpty)
	}
}

func TestStationNeighborsTableMarkExcluded(t *testing.T) {
	snt := NewStationNeighborsTable()
	snt.Add("c1", "s1", "Station One", "geo")
	snt.Add("c1", "s2", "Station Two", "geo")

	snt.MarkExcluded("c1", "s1")

	neighbors := snt.GetByCity("c1")
	if len(neighbors) != 2 {
		t.Errorf("expected 2 neighbors, got %d", len(neighbors))
	}

	if !neighbors[0].IsExcluded {
		t.Error("expected s1 to be excluded")
	}
	if neighbors[1].IsExcluded {
		t.Error("expected s2 to not be excluded")
	}
}

func TestStationNeighborsTableGetNonExcluded(t *testing.T) {
	snt := NewStationNeighborsTable()
	snt.Add("c1", "s1", "Station One", "geo")
	snt.Add("c1", "s2", "Station Two", "geo")
	snt.MarkExcluded("c1", "s1")

	nonExcluded := snt.GetNonExcluded("c1")
	if len(nonExcluded) != 1 {
		t.Errorf("expected 1 non-excluded neighbor, got %d", len(nonExcluded))
	}

	if nonExcluded[0].StationID != "s2" {
		t.Errorf("expected 's2' as non-excluded, got '%s'", nonExcluded[0].StationID)
	}
}