#!/usr/bin/env python3
"""Quick test to verify word length distribution in generated puzzles."""

from wrdler.generator import generate_puzzle
from wrdler.word_loader import load_word_list

def test_word_distribution():
    """Test that puzzles have 2 four-letter, 2 five-letter, and 2 six-letter words."""
    print("Testing word distribution in generated puzzles...")

    # Load word list
    words_by_len = load_word_list()
    print(f"Loaded words: {len(words_by_len[4])} 4-letter, {len(words_by_len[5])} 5-letter, {len(words_by_len[6])} 6-letter")

    # Generate 5 test puzzles
    for i in range(5):
        puzzle = generate_puzzle(grid_rows=6, grid_cols=8, words_by_len=words_by_len)

        # Count word lengths
        length_counts = {4: 0, 5: 0, 6: 0}
        for word in puzzle.words:
            length = len(word.text)
            if length in length_counts:
                length_counts[length] += 1

        # Verify distribution
        assert length_counts[4] == 2, f"Puzzle {i+1}: Expected 2 four-letter words, got {length_counts[4]}"
        assert length_counts[5] == 2, f"Puzzle {i+1}: Expected 2 five-letter words, got {length_counts[5]}"
        assert length_counts[6] == 2, f"Puzzle {i+1}: Expected 2 six-letter words, got {length_counts[6]}"

        # Print puzzle info
        words_str = ", ".join([f"{w.text}({len(w.text)})" for w in puzzle.words])
        print(f"✓ Puzzle {i+1}: {words_str}")

    print("\n✅ All tests passed! Word distribution is correct.")

if __name__ == "__main__":
    test_word_distribution()