package trie

import (
	"testing"
)

func TestNew(t *testing.T) {
	root := New()
	if root == nil {
		t.Error("New() returned nil")
	}
	if root.children == nil {
		t.Error("New() root node has nil children map")
	}
	if root.isEnd {
		t.Error("New() root node should not be marked as end")
	}
}

func TestAddAndCount(t *testing.T) {
	tests := []struct {
		name     string
		words    []string
		prefix   string
		expected int
	}{
		{
			name:     "empty trie",
			words:    []string{},
			prefix:   "test",
			expected: 0,
		},
		{
			name:     "single word",
			words:    []string{"hello"},
			prefix:   "hel",
			expected: 1,
		},
		{
			name:     "multiple words same prefix",
			words:    []string{"hello", "help", "hell"},
			prefix:   "hel",
			expected: 3,
		},
		{
			name:     "prefix not found",
			words:    []string{"hello", "help", "hell"},
			prefix:   "xyz",
			expected: 0,
		},
		{
			name:     "empty prefix",
			words:    []string{"hello", "help", "hell"},
			prefix:   "",
			expected: 3,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			root := New()
			for _, word := range tt.words {
				root.Add(word)
			}
			if got := root.Count(tt.prefix); got != tt.expected {
				t.Errorf("Count(%q) = %d, want %d", tt.prefix, got, tt.expected)
			}
		})
	}
}

func TestTopPrefixes(t *testing.T) {
	tests := []struct {
		name     string
		words    []string
		expected []PrefixCount
	}{
		{
			name:     "empty trie",
			words:    []string{},
			expected: []PrefixCount{},
		},
		{
			name:  "single word",
			words: []string{"hello"},
			expected: []PrefixCount{
				{Prefix: "", Count: 1},
				{Prefix: "h", Count: 1},
				{Prefix: "he", Count: 1},
				{Prefix: "hel", Count: 1},
				{Prefix: "hell", Count: 1},
				{Prefix: "hello", Count: 1},
			},
		},
		{
			name:  "multiple words",
			words: []string{"hello", "help", "hell", "helicopter"},
			expected: []PrefixCount{
				{Prefix: "", Count: 4},
				{Prefix: "h", Count: 4},
				{Prefix: "he", Count: 4},
				{Prefix: "hel", Count: 4},
				{Prefix: "hell", Count: 2},
				{Prefix: "hello", Count: 1},
				{Prefix: "help", Count: 1},
				{Prefix: "heli", Count: 1},
				{Prefix: "helic", Count: 1},
				{Prefix: "helico", Count: 1},
				{Prefix: "helicop", Count: 1},
				{Prefix: "helicopt", Count: 1},
				{Prefix: "helicopte", Count: 1},
				{Prefix: "helicopter", Count: 1},
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			root := New()
			for _, word := range tt.words {
				root.Add(word)
			}
			got := TopPrefixes(root)
			if len(got) != len(tt.expected) {
				t.Errorf("TopPrefixes() returned %d results, want %d", len(got), len(tt.expected))
				return
			}
			for i, want := range tt.expected {
				if i >= len(got) {
					break
				}
				if got[i].Prefix != want.Prefix || got[i].Count != want.Count {
					t.Errorf("TopPrefixes()[%d] = {Prefix: %q, Count: %d}, want {Prefix: %q, Count: %d}",
						i, got[i].Prefix, got[i].Count, want.Prefix, want.Count)
				}
			}
		})
	}
}

func TestTopPrefixesSorting(t *testing.T) {
	root := New()
	words := []string{"a", "ab", "abc", "abcd", "abcde"}
	for _, word := range words {
		root.Add(word)
	}

	prefixes := TopPrefixes(root)

	// Verify sorting by count (descending)
	for i := 1; i < len(prefixes); i++ {
		if prefixes[i].Count > prefixes[i-1].Count {
			t.Errorf("Prefixes not sorted by count: %v", prefixes)
		}
	}

	// Verify alphabetical sorting for equal counts
	for i := 1; i < len(prefixes); i++ {
		if prefixes[i].Count == prefixes[i-1].Count && prefixes[i].Prefix < prefixes[i-1].Prefix {
			t.Errorf("Prefixes with equal counts not sorted alphabetically: %v", prefixes)
		}
	}
}