"""
Functions for declaring constraints.

This module provides functions for declaring constraints on data models including:
- Uniqueness constraints
- Exclusive/anyof/oneof constraints for relationships
"""
from __future__ import annotations

from ..frontend.base import FieldRef, Fragment, Library, Expression, Field, TupleVariable, Variable, Concept, Ref, MetaRef
from ..frontend import core
from relationalai.util.docutils import include_in_docs

__include_in_docs__ = True

# the front-end library object
library = Library("constraints")

#------------------------------------------------------
# Constraint Relations
#------------------------------------------------------

_unique = library.Relation("unique", [Field.input("values", core.Any, is_list=True)])
_unique_fields = library.Relation("unique_fields", [Field.input("fields", core.Field, is_list=True)])
_exclusive = library.Relation("exclusive", [Field.input("concepts", core.Relation, is_list=True)])
_anyof =library.Relation("anyof", [Field.input("concepts", core.Relation, is_list=True)])

#------------------------------------------------------
# API
#------------------------------------------------------

@include_in_docs
def unique(*args: Variable) -> Expression:
    """
    Declare a uniqueness constraint on variables.

    Parameters
    ----------
    *args: Variable
        Variables representing values or fields that must be unique.

    Returns
    -------
    Expression
        An `Expression` computing the uniqueness constraint.

    """
    if all(isinstance(arg, FieldRef) for arg in args):
        fields = [MetaRef(arg._resolved) for arg in args] # type: ignore - we know these are FieldRefs from the all above
        return _unique_fields(TupleVariable(fields))
    return _unique(TupleVariable(args))

@include_in_docs
def exclusive(*args: Concept|Ref) -> Expression:
    """
    Declare that the given concepts are mutually exclusive.

    Parameters
    ----------
    *args: Concept | Ref
        Concepts or Refs that should be mutually exclusive.

    Returns
    -------
    Expression
        An `Expression` computing the exclusivity constraint.
    """
    concepts = [arg._to_concept() for arg in args]
    return _exclusive(TupleVariable(concepts))

@include_in_docs
def anyof(*args: Concept|Ref) -> Expression:
    """
    Declare that at least one of the given concepts must be true.

    Parameters
    ----------
    *args: Concept | Ref
        Concepts or Refs, at least one of which must be satisfied.

    Returns
    -------
    Expression
        An `Expression` computing the anyof constraint.
    """
    concepts = [arg._to_concept() for arg in args]
    return _anyof(TupleVariable(concepts))

@include_in_docs
def oneof(*args: Concept | Ref) -> Fragment:
    """
    Declare that exactly one of the given concepts must be true.

    This combines `anyof` (at least one) and `exclusive` (at most one) constraints.

    Parameters
    ----------
    *args: Concept | Ref
        Concepts or Refs, exactly one of which must be satisfied.

    Returns
    -------
    Fragment
        A `Fragment` representing the oneof constraint.
    """
    model = args[0]._model
    return model.where(anyof(*args), exclusive(*args))
