"""Mini Sudoku (4x4) puzzle generator"""

import copy
from dataclasses import dataclass
from random import Random
from typing import Any, Optional

from ..coaching import BaseCurriculum, RangeAttributeDefinition
from ..factory import ProceduralDataset, register_dataset

DATASET_NAME = "mini_sudoku"


@dataclass
class MiniSudokuConfig:
    """Configuration for 4x4 sudoku puzzle generation"""

    min_empty: int = (
        8  # Minimum number of empty cells. Occasionally this can be violated, if removing more cells would break the puzzle's uniqueness.
    )
    max_empty: int = 12  # Maximum number of empty cells
    seed: Optional[int] = None
    size: int = 500  # Virtual dataset size

    def validate(self):
        """Validate configuration parameters"""
        # More than 12 empty cells is incompatible with a unique solution
        assert 0 <= self.min_empty <= 12, "min_empty must be between 0 and 12"
        assert self.min_empty <= self.max_empty <= 12, "max_empty must be between min_empty and 12"


class MiniSudokuDataset(ProceduralDataset):
    """Generates 4x4 sudoku puzzles with configurable difficulty"""

    def __init__(self, config: MiniSudokuConfig):
        super().__init__(config=config, seed=config.seed, size=config.size)

    def __len__(self) -> int:
        return self.config.size

    def __iter__(self):
        self._current_idx = 0
        return self

    def __next__(self):
        if self._current_idx >= self.config.size:
            raise StopIteration
        item = self[self._current_idx]
        self._current_idx += 1
        return item

    def _is_valid(self, board: list[list[int]], row: int, col: int, num: int) -> bool:
        """Check if number can be placed at position"""
        # Check row
        if num in board[row]:
            return False

        # Check column
        if num in [board[i][col] for i in range(4)]:
            return False

        # Check 2x2 box
        box_row, box_col = 2 * (row // 2), 2 * (col // 2)
        for i in range(box_row, box_row + 2):
            for j in range(box_col, box_col + 2):
                if board[i][j] == num:
                    return False
        return True

    def _solve(self, board: list[list[int]]) -> bool:
        """Solve mini sudoku using backtracking"""
        empty = self._find_empty(board)
        if not empty:
            return True

        row, col = empty
        for num in range(1, 5):
            if self._is_valid(board, row, col, num):
                board[row][col] = num
                if self._solve(board):
                    return True
                board[row][col] = 0
        return False

    def _find_empty(self, board: list[list[int]]) -> Optional[tuple[int, int]]:
        """Find an empty cell"""
        for i in range(4):
            for j in range(4):
                if board[i][j] == 0:
                    return (i, j)
        return None

    def _generate_solved_board(self, rng: Random) -> list[list[int]]:
        """Generate a complete solved mini sudoku board"""
        board = [[0] * 4 for _ in range(4)]

        # Try multiple times to generate a valid board
        max_attempts = 100
        for _ in range(max_attempts):
            # Start fresh
            for i in range(4):
                for j in range(4):
                    board[i][j] = 0

            # Fill diagonal boxes first (they are independent)
            for i in range(0, 4, 2):
                nums = list(range(1, 5))
                rng.shuffle(nums)
                pos = 0
                for r in range(i, i + 2):
                    for c in range(i, i + 2):
                        board[r][c] = nums[pos]
                        pos += 1

            # Try to solve the rest
            if self._solve(board):
                return board

        raise RuntimeError("Failed to generate valid mini sudoku board")

    def _count_solutions(self, board: list[list[int]], limit: int = 2) -> int:
        """Count the number of solutions for a given board"""

        def _count_solutions_helper(board: list[list[int]]) -> int:
            empty = self._find_empty(board)
            if not empty:
                return 1

            row, col = empty
            count = 0
            for num in range(1, 5):
                if self._is_valid(board, row, col, num):
                    board[row][col] = num
                    count += _count_solutions_helper(board)
                    if count >= limit:
                        return count
                    board[row][col] = 0
            return count

        return _count_solutions_helper(board)

    def _create_puzzle(self, solved_board: list[list[int]], num_empty: int, rng: Random) -> list[list[int]]:
        """Create puzzle by removing numbers from solved board"""
        puzzle = [row[:] for row in solved_board]
        cells = [(i, j) for i in range(4) for j in range(4)]
        rng.shuffle(cells)
        num_removed = 0

        for i, j in cells:
            saved = puzzle[i][j]
            puzzle[i][j] = 0
            puzzle_copy = copy.deepcopy(puzzle)
            # Check if removing this clue breaks uniqueness
            if self._count_solutions(puzzle_copy) > 1:
                puzzle[i][j] = saved
            else:
                num_removed += 1
                if num_removed == num_empty:
                    break

        return puzzle

    def _board_to_string(self, board: list[list[int]]) -> str:
        """Convert board to string representation"""
        return "\n".join(" ".join(str(x) if x != 0 else "_" for x in row) for row in board)

    def __getitem__(self, idx: int) -> dict:
        """Generate a single mini sudoku puzzle"""
        rng = Random(self.seed + idx)

        # Generate solved board
        solved_board = self._generate_solved_board(rng)

        # Create puzzle by removing numbers
        num_empty = rng.randint(self.config.min_empty, self.config.max_empty)
        puzzle = self._create_puzzle(solved_board, num_empty, rng)

        # Update the num_empty to be used in the metadata if we couldn't remove as many as we wanted
        num_empty = sum(1 for row in puzzle for x in row if x == 0)

        # Format as strings
        puzzle_str = self._board_to_string(puzzle)
        solution_str = self._board_to_string(solved_board)

        question = (
            "In 4x4 Mini Sudoku:\n"
            "- Each row must contain each number from 1-4 exactly once\n"
            "- Each column must contain each number 1-4 exactly once\n"
            "- Each 2x2 subgrid must contain each number 1-4 exactly once\n"
            f"Solve this 4x4 Mini Sudoku puzzle:\n{puzzle_str}\n"
            "Format your response as the puzzle above, with spaces separating each number within a row, and newlines separating rows.\n"
        )

        return {
            "question": question,
            "answer": solution_str,
            "metadata": {
                "source_dataset": DATASET_NAME,
                "source_index": idx,
                "puzzle": puzzle,
                "solution": solved_board,
                "num_empty": num_empty,
                "difficulty": {
                    "empty": (self.config.min_empty, self.config.max_empty),
                },
            },
        }

    def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float:
        if not isinstance(answer, str) or len(answer) == 0:
            return 0.0

        oracle_answer = entry["answer"]
        metadata = entry["metadata"]
        solution: list[list[int]] = metadata["solution"]
        board_size: int = len(solution[0])

        # 1. match answer without trailing whitespaces
        answer_stripped = "\n".join(l.rstrip() for l in answer.split("\n"))
        oracle_answer_stripped = "\n".join(l.rstrip() for l in oracle_answer.split("\n"))

        if answer_stripped == oracle_answer_stripped:
            reward = 1.0
        else:
            # 2. accept answers with correct numeric sequence (ignoring non-numeric characters)
            row = 0
            num_matching = 0
            for ln in answer.split("\n"):
                if row >= len(solution):
                    break
                numbers = [int(c) for c in ln if c in "123456789"]
                if len(numbers) != board_size:
                    continue  # ignore lines without numbers
                for a, b in zip(solution[row], numbers):
                    if a == b:
                        num_matching += 1
                row += 1

            reward = num_matching / (board_size * board_size)
            reward *= 0.9  # penalty for not using standard format

        if len(answer) > len(oracle_answer):
            reward *= len(oracle_answer) / len(answer)  # penalty for additional length
        return reward


class MiniSudokuCurriculum(BaseCurriculum):
    def __init__(self):
        super().__init__(MiniSudokuCurriculum.__name__, MiniSudokuConfig)

        # Define attributes
        self._define_attributes(
            RangeAttributeDefinition(
                name="empty",
                levels=[4, 6, 8, 10, 12],
                description="Number of empty cells in the puzzle",
                lower_field_name="min_empty",
                upper_field_name="max_empty",
                ensure_interval=True,
            )
        )


register_dataset(DATASET_NAME, MiniSudokuDataset, MiniSudokuConfig, MiniSudokuCurriculum)
