"""
norvig.csp – AIMA Chapter 6
=============================
Uses the ``constraint`` library (python-constraint2) as the solver backend.
Provides a thin AIMA-flavoured wrapper and re-exports the library's
AllDifferentConstraint / InSetConstraint for common patterns.

    from ml_ai_library.norvig.csp import (
        CSP,                    # AIMA-style wrapper over constraint.Problem
        AllDifferentConstraint, # from constraint library
        solve,                  # return first solution or None
        solve_all,              # return all solutions
    )

Example – Map colouring
-----------------------
>>> from ml_ai_library.norvig.csp import CSP, AllDifferentConstraint, solve
>>> csp = CSP(
...     variables=["WA","NT","SA","Q","NSW","V"],
...     domains={"WA":["R","G","B"], ...},
... )
>>> csp.add_constraint(AllDifferentConstraint(), ["WA","NT"])
>>> csp.add_constraint(AllDifferentConstraint(), ["WA","SA"])
>>> result = solve(csp)
"""

from __future__ import annotations
from typing import Any, Callable, Dict, List, Optional

# python-constraint2 library as the CSP engine --------------------------------
from constraint import (
    Problem as _ConstraintProblem,
    AllDifferentConstraint,
    InSetConstraint,
    FunctionConstraint,
)


class CSP:
    """
    AIMA-style CSP wrapper over the ``constraint`` library (AIMA §6.1).

    Parameters
    ----------
    variables : list of variable names
    domains   : dict {var: list_of_values}
    """

    def __init__(self, variables: List[Any], domains: Dict[Any, List[Any]]):
        self._problem = _ConstraintProblem()
        for var in variables:
            self._problem.addVariable(var, list(domains[var]))

    def add_constraint(self, constraint, variables: List[Any]):
        """
        Add a constraint.  Pass a constraint object (AllDifferentConstraint,
        InSetConstraint) or a callable (X, Y) -> bool.
        """
        if callable(constraint) and not hasattr(constraint, "__call_with_variables"):
            self._problem.addConstraint(FunctionConstraint(constraint), variables)
        else:
            self._problem.addConstraint(constraint, variables)

    def add_binary_constraint(self, fn: Callable, var1, var2):
        """Convenience: add a binary constraint fn(x, y) -> bool."""
        self._problem.addConstraint(fn, [var1, var2])

    @property
    def _raw(self) -> _ConstraintProblem:
        return self._problem


def solve(csp: CSP) -> Optional[Dict]:
    """Return one solution dict or None (backtracking + AC via constraint lib)."""
    return csp._raw.getSolution()


def solve_all(csp: CSP) -> List[Dict]:
    """Return all solution dicts."""
    return csp._raw.getSolutions()


__all__ = [
    "CSP", "solve", "solve_all",
    "AllDifferentConstraint", "InSetConstraint", "FunctionConstraint",
]
