from __future__ import annotations

import random
import hashlib  # NEW
from typing import Dict, List, Optional, Union  # UPDATED
from .word_loader import load_word_list
from .models import Coord, Word, Puzzle


def _fits_and_free(cells: List[Coord], used: set[Coord], size: int) -> bool:
    """Legacy function for square grids (deprecated)."""
    for c in cells:
        if not c.in_bounds(size) or c in used:
            return False
    return True


def _fits_and_free_rect(cells: List[Coord], used: set[Coord], rows: int, cols: int) -> bool:
    """Check if cells fit in rectangular grid and are free."""
    for c in cells:
        if not (0 <= c.x < rows and 0 <= c.y < cols) or c in used:
            return False
    return True


def _build_cells(start: Coord, length: int, direction: str) -> List[Coord]:
    if direction == "H":
        return [Coord(start.x, start.y + i) for i in range(length)]
    else:
        return [Coord(start.x + i, start.y) for i in range(length)]


def _chebyshev_distance(a: Coord, b: Coord) -> int:
    return max(abs(a.x - b.x), abs(a.y - b.y))


def _seed_from_id(puzzle_id: str) -> int:
    """Derive a deterministic 64-bit seed from a string id."""
    h = hashlib.sha256(puzzle_id.encode("utf-8")).digest()
    return int.from_bytes(h[:8], "big", signed=False)


def generate_puzzle(
    grid_rows: int = 6,
    grid_cols: int = 8,
    grid_size: Optional[int] = None,  # Legacy parameter for backward compatibility
    words_by_len: Optional[Dict[int, List[str]]] = None,
    seed: Optional[Union[int, str]] = None,
    max_attempts: int = 5000,
    spacer: int = 1,
    puzzle_id: Optional[str] = None,
    _retry: int = 0,
    target_words: Optional[List[str]] = None,
) -> Puzzle:
    """
    Generate a Wrdler puzzle with 6 horizontal words (one per row).

    Wrdler Specifications:
    - 6 rows × 8 columns grid
    - 6 words total (one per row)
    - **Word distribution:** 2 four-letter words, 2 five-letter words, 2 six-letter words
    - All words horizontal only
    - No vertical words, no overlaps

    Parameters
    - grid_rows: number of rows (default 6)
    - grid_cols: number of columns (default 8)
    - grid_size: legacy parameter for square grids (deprecated, use grid_rows/grid_cols)
    - words_by_len: preloaded word pools by length
    - seed: optional RNG seed for reproducibility
    - max_attempts: cap for word selection attempts
    - spacer: separation constraint (0=touching, 1=1 space, 2=2 spaces)
    - puzzle_id: deterministic puzzle identifier
    - _retry: internal retry counter for deterministic regeneration
    - target_words: specific list of 6 words to use (for shared games)

    Returns
    - Puzzle object with 6 horizontal words

    Determinism:
    - If puzzle_id is provided, it's used to derive the RNG seed
    - Retries use f"{puzzle_id}:{_retry}" for deterministic regeneration
    - Else if seed is provided (int or str), it's used directly
    - Else RNG is non-deterministic
    """
    # Handle legacy grid_size parameter
    if grid_size is not None:
        grid_rows = grid_size
        grid_cols = grid_size
    # Compute deterministic seed if requested
    if puzzle_id is not None:
        seed_val = _seed_from_id(f"{puzzle_id}:{_retry}")
    elif isinstance(seed, str):
        seed_val = _seed_from_id(f"{seed}:{_retry}")
    elif isinstance(seed, int):
        seed_val = seed + _retry
    else:
        seed_val = None

    rng = random.Random(seed_val) if seed_val is not None else random.Random()

    # Validate grid dimensions for Wrdler
    if grid_rows != 6:
        raise ValueError(f"Wrdler requires exactly 6 rows, got {grid_rows}")
    if grid_cols < 3:
        raise ValueError(f"Grid must have at least 3 columns, got {grid_cols}")

    # If target_words is provided, use those specific words
    if target_words:
        if len(target_words) != 6:
            raise ValueError(f"target_words must contain exactly 6 words, got {len(target_words)}")
        # Validate word lengths fit in grid
        for word in target_words:
            if len(word) > grid_cols:
                raise ValueError(f"Word '{word}' ({len(word)} letters) too long for {grid_cols} columns")
            if len(word) < 3:
                raise ValueError(f"Word '{word}' too short (minimum 3 letters)")
        selected_words = [w.upper() for w in target_words]
    else:
        # Normal random word selection - select 6 words with required distribution
        # Wrdler requires: 2 four-letter words, 2 five-letter words, 2 six-letter words
        words_by_len = words_by_len or load_word_list()

        # Validate we have enough words in each required length
        required_lengths = [4, 5, 6]
        for length in required_lengths:
            if length not in words_by_len or len(words_by_len[length]) < 2:
                raise RuntimeError(
                    f"Insufficient {length}-letter words (need at least 2, found {len(words_by_len.get(length, []))})"
                )

        # Select 2 words from each length pool
        selected_words: List[str] = []
        for length in required_lengths:
            pool = list(dict.fromkeys(words_by_len[length]))  # Deduplicate
            rng.shuffle(pool)
            selected_words.extend(pool[:2])  # Take 2 words from this length

        # Shuffle the selected 6 words to randomize their order in the grid
        rng.shuffle(selected_words)

    # Wrdler placement algorithm: one word per row, horizontal only
    # Shuffle row order for variety
    row_indices = list(range(grid_rows))
    rng.shuffle(row_indices)

    placed: List[Word] = []

    for i, word_text in enumerate(selected_words):
        row_idx = row_indices[i]
        word_len = len(word_text)

        # Calculate valid starting positions for this word in its row
        # Word must fit within grid_cols
        if word_len > grid_cols:
            raise RuntimeError(f"Word '{word_text}' ({word_len} letters) doesn't fit in {grid_cols} columns")

        max_start_col = grid_cols - word_len

        # Randomly choose starting column
        if max_start_col >= 0:
            start_col = rng.randint(0, max_start_col)
        else:
            start_col = 0

        # Create word (always horizontal)
        word = Word(
            text=word_text,
            start=Coord(row_idx, start_col),
            direction="H"
        )
        placed.append(word)

    # Create puzzle
    puzzle = Puzzle(
        words=placed,
        spacer=spacer,
        grid_rows=grid_rows,
        grid_cols=grid_cols
    )

    # Validate
    try:
        validate_puzzle(puzzle, grid_rows=grid_rows, grid_cols=grid_cols)
    except AssertionError as e:
        # Deterministic retry on validation failure
        if _retry >= 10:
            raise RuntimeError(f"Puzzle generation failed after {_retry} retries: {e}")

        return generate_puzzle(
            grid_rows=grid_rows,
            grid_cols=grid_cols,
            words_by_len=words_by_len,
            seed=seed,
            max_attempts=max_attempts,
            spacer=spacer,
            puzzle_id=puzzle_id,
            _retry=_retry + 1,
            target_words=target_words,
        )

    return puzzle


def validate_puzzle(
    puzzle: Puzzle,
    grid_rows: int = 6,
    grid_cols: int = 8,
    grid_size: Optional[int] = None  # Legacy parameter
) -> None:
    """
    Validate Wrdler puzzle constraints.

    Checks:
    1. Exactly 6 words
    2. All words are horizontal
    3. One word per row
    4. All cells within grid bounds
    5. No overlapping cells
    6. Word length distribution: exactly 2 four-letter, 2 five-letter, 2 six-letter words
    """
    # Handle legacy grid_size parameter
    if grid_size is not None:
        grid_rows = grid_size
        grid_cols = grid_size

    # 1. Check word count
    if len(puzzle.words) != 6:
        raise AssertionError(f"Expected exactly 6 words, got {len(puzzle.words)}")

    # 2. Check all horizontal
    for w in puzzle.words:
        if w.direction != "H":
            raise AssertionError(f"Word '{w.text}' is not horizontal (direction={w.direction})")

    # 3. Check one word per row
    rows_used = [w.start.x for w in puzzle.words]
    if len(set(rows_used)) != 6:
        raise AssertionError(f"Must have one word per row, got {len(set(rows_used))} unique rows")

    # Ensure all 6 rows are used
    if set(rows_used) != set(range(6)):
        raise AssertionError(f"All 6 rows must be used, got rows: {sorted(set(rows_used))}")

    # 4. Check bounds and overlaps
    seen: set[Coord] = set()
    word_lengths: list[int] = []
    for w in puzzle.words:
        word_lengths.append(len(w.text))

        # Check word length
        if len(w.text) < 3:
            raise AssertionError(f"Word '{w.text}' too short (< 3 letters)")
        if len(w.text) > grid_cols:
            raise AssertionError(f"Word '{w.text}' too long ({len(w.text)} > {grid_cols})")

        # Check all cells in bounds
        for c in w.cells:
            if not (0 <= c.x < grid_rows and 0 <= c.y < grid_cols):
                raise AssertionError(f"Cell {c} out of bounds (grid is {grid_rows}×{grid_cols})")

            # Check no overlaps (should be impossible with one-per-row, but verify)
            if c in seen:
                raise AssertionError(f"Overlapping cell detected: {c}")
            seen.add(c)

    # 5. Check word length distribution (Wrdler requirement)
    # Must have exactly: 2 four-letter, 2 five-letter, 2 six-letter words
    length_counts = {4: 0, 5: 0, 6: 0}
    for length in word_lengths:
        if length in length_counts:
            length_counts[length] += 1
        else:
            raise AssertionError(f"Invalid word length {length} (must be 4, 5, or 6)")

    if length_counts[4] != 2:
        raise AssertionError(f"Must have exactly 2 four-letter words, got {length_counts[4]}")
    if length_counts[5] != 2:
        raise AssertionError(f"Must have exactly 2 five-letter words, got {length_counts[5]}")
    if length_counts[6] != 2:
        raise AssertionError(f"Must have exactly 2 six-letter words, got {length_counts[6]}")


    # Note: Spacer rules not needed for Wrdler since words are on different rows
    # They cannot touch each other (minimum 1 row separation)

def sort_word_file(filepath: str) -> List[str]:
    """
    Reads a word list file, skips header/comment lines, and returns words sorted
    by length (ascending), then alphabetically within each length group.
    """
    with open(filepath, "r", encoding="utf-8") as f:
        lines = f.readlines()
    # Skip header/comment lines
    words = [line.strip() for line in lines if line.strip() and not line.strip().startswith("#")]
    # Sort by length, then alphabetically
    sorted_words = sorted(words, key=lambda w: (len(w), w))
    return sorted_words