from typing import Optional, Sequence
from ....metamodel.metamodel import (
    Require, Logical, Var, Relation, Lookup, ScalarType, Literal, Field
)
from ....metamodel.builtins import builtins

"""
Helper functions for converting `Require` nodes with unique constraints to functional
dependencies. The main functionalities provided are:
    1. Check whether a `Require` node is a valid unique constraint representation
    2. Represent the uniqueness constraint as a functional dependency
    3. Check if the functional dependency is structural i.e., can be represented with
       `@function(k)` annotation on a single relation.

=========================== Structure of unique constraints ================================
A `Require` node represents a _unique constraint_ if it meets the following criteria:
 * the `Require` node's `domain` is an empty `Logical` node
 * the `Require` node's `check` is a `Logical` node containing one or more `unique_fields`
   `Lookup` tasks
 * each `unique_fields` lookup has precisely one argument, which is a tuple of `Field` objects
 * all fields in a `unique_fields` lookup belong to the same relation
 * the fields represent the unique columns in the relation

Note: In the v1 metamodel, the guard atoms (relation lookups) are implicit in the fields.
The fields contain references to their parent relation via the `_relation` attribute.
============================================================================================

We use the following unique constraint as the running example (v1 metamodel structure).

```
Require
  domain:
    Logical([])
  check:
    Logical([
      unique_fields((first_name_field, last_name_field))
    ])
```

where `first_name_field` and `last_name_field` are `Field` objects that both belong to
a relation like `Person(person::Person, firstname::String, lastname::String)`.

This is equivalent to the v0 representation:
```
Require
  domain:
    Logical([])
  checks:
    Check(
      check:
        Logical([
          Person(person::Person, firstname::String, lastname::String),
          unique((firstname::String, lastname::String))
        ])
    )
```

=========================== Semantics of unique constraints ================================
A unique constraint states that the columns declared in the `unique` predicate must be
unique in the result of the conjunctive query consisting of all remaining predicates.
============================================================================================

In the running example, the conjunctive query computes a table with 3 columns, the person id
`person::Person`, the first name `firstname::String`, and the last name `lastname::String`.
The uniqueness predicate `unique((firstname::String, lastname::String))` states that no person
can have more than a single combination of first and last name.

The unique constraint in the running example above corresponds to the following functional
dependency.

```
Person(x) ∧ first_name(x, y) ∧ last_name(x, z): {y, z} -> {x}
```

------------------------------ Redundant Type Atoms ----------------------------------------
At the time of writing, PyRel does not yet remove redundant unary atoms. For instance, in
the running example, the atom `Person(person::Person)` is redundant because the type of the
`person` variable is specified in the other two atoms `first_name` and `last_name`.
Consequently, we identify redundant atoms and remove them from the definition of the
corresponding functional dependency.

Formally, a _guard_ atom is any `Lookup` node whose relation name is not `unique`. Now, a
unary guard atom `T(x::T)` is _redundant_ if the uniqueness constraint has a non-unary guard
atom `R(...,x::T,...)`.

================================ Normalized FDs ============================================
Now, the _(normalized)_ functional dependency_ corresponding to a unique constraint is an
object of the form `φ: X → Y`, where :
 1. `φ` is the set of all non-redundant guard atoms.
 2. `X` is the set of variables used in the `unique` atom
 3. `Y` is the set of all other variables used in the constraint
============================================================================================

The normalized functional dependency corresponding to the unique constraints from the running
example is :
```
first_name(person::Person, firstname::String) ∧ last_name(person::Person, lastname::String): {firstname:String, lastname:String} -> {person:Person}
```
Note that the unary atom `Person(person::Person)` is redundant and thus omitted from the
decomposition.

Some simple functional dependencies can, however, be expressed simply with `@function(k)`
attribute of a single relation.  Specifically, a functional dependency `φ: X → Y` is
_structural_ if φ consists of a single atom `R(x1,...,xm,y1,...,yk)` and `X = {x1,...,xm}`.
"""

#
# Processes a single `unique_fields` Lookup and extracts the guard and variables.
# Returns `None` if the lookup is not a valid unique_fields constraint.
# If valid, returns a tuple `(all_vars, unique_vars, guard)`, where:
# - `all_vars` is the list of all variables used in the constraint
# - `unique_vars` is the list of variables that are declared unique
# - `guard` is the list of guard `Lookup` atoms (constructed from the fields)
#
def _split_unique_fields_lookup(lookup: Lookup) -> Optional[tuple[list[Var], list[Var], list[Lookup]]]:
    if lookup.relation.name != builtins.constraints.unique_fields.name:
        return None

    # Extract fields from unique_fields((field1, field2, ...))
    if len(lookup.args) != 1:
        return None
    fields_arg = lookup.args[0]
    if not isinstance(fields_arg, tuple):
        return None
    if len(fields_arg) == 0:
        return None

    # All fields should be Field objects
    if not all(isinstance(f, Field) for f in fields_arg):
        return None

    # Now we know all elements are Field objects
    fields: list[Field] = list(fields_arg)  # type: ignore

    # Get the relation from the first field
    # All fields should belong to the same relation (this is the guard relation)
    first_field = fields[0]
    relation = first_field._relation
    if relation is None:
        return None

    # Verify all fields belong to the same relation
    for f in fields:
        if f._relation != relation:
            return None

    # Create variables for all fields in the relation
    all_vars: list[Var] = []
    for field in relation.fields:
        var = Var(type=field.type, name=field.name)
        all_vars.append(var)

    # The unique_vars are the variables corresponding to the fields in unique_fields
    unique_vars: list[Var] = []
    for field in fields:
        # Find the index of this field in the relation's fields
        try:
            idx = relation.fields.index(field)
            unique_vars.append(all_vars[idx])
        except ValueError:
            return None

    # Create the guard lookup: Relation(all_vars)
    guard_lookup = Lookup(relation, tuple(all_vars))
    guard = [guard_lookup]

    return list(all_vars), list(unique_vars), guard

def get_unique_fields_lookups(node: Require) -> list[Lookup]:
    """
    Extracts all `unique_fields` lookups from a `Require` node.
    Returns an empty list if the node is not a valid unique constraint.
    """
    assert isinstance(node.check, Logical)
    lookups = []
    for task in node.check.body:
        if isinstance(task, Lookup) and task.relation.name == builtins.constraints.unique_fields.name:
            lookups.append(task)

    return lookups

class FunctionalDependency:
    """
    Represents a functional dependency of the form `φ: X -> Y`, where
     - `φ` is a set of `Lookup` atoms
     - `X` and `Y` are disjoint and covering sets of variables used in `φ`
    """
    def __init__(self, guard: Sequence[Lookup], keys: Sequence[Var], values: Sequence[Var]):
        self.guard = tuple(guard)
        self.keys = tuple(keys)
        self.values = tuple(values)
        assert set(self.keys).isdisjoint(set(self.values)), "Keys and values must be disjoint"

        # for structural fd check
        self._is_structural:bool = False
        self._structural_relation:Optional[Relation] = None
        self._structural_rank:Optional[int] = None

        self._determine_is_structural()

        # compute canonical string representation of the fd
        self._canonical_str = self._compute_canonical_str()

    # A functional dependency `φ: X → Y` is _k-functional_ if `φ` consists of a single atom
    # `R(x1,...,xm,y1,...,yk)` and `X = {x1,...,xm}`. Not all functional dependencies are
    # k-functional. For instance, `R(x, y, z): {y, z} → {x}` cannot be expressed with
    # `@function`. neither can `R(x, y) ∧ P(x, z) : {x} → {y, z}`.
    def _determine_is_structural(self):
        if len(self.guard) != 1:
            self._is_structural = False
            return
        atom = next(iter(self.guard))
        atom_vars = atom.args
        if len(atom_vars) <= len(self.keys): # @function(0) provides no information
            self._is_structural = False
            return
        prefix_vars = atom_vars[:len(self.keys)]
        if set(prefix_vars) != set(self.keys):
            self._is_structural = False
            return
        self._is_structural = True
        self._structural_relation = atom.relation
        self._structural_rank = len(atom_vars) - len(self.keys)

    @property
    def is_structural(self) -> bool:
        """
        Whether the functional dependency is functional, i.e., can be represented
        with `@function(k)` annotation on a single relation.
        """
        return self._is_structural

    @property
    def structural_relation(self) -> Relation:
        """
        The structural relation of a functional dependency. Raises ValueError if the functional
        dependency is not structural.
        """
        if not self._is_structural:
            raise ValueError("Functional dependency is not structural")
        assert self._structural_relation is not None
        return self._structural_relation

    @property
    def structural_rank(self) -> int:
        """
        The structural rank k of k-structural fd. Raises ValueError if the structural
        dependency is not k-structural.
        """
        if not self._is_structural:
            raise ValueError("Functional dependency is not structural")
        assert self._structural_rank is not None
        return self._structural_rank

    def __str__(self) -> str:
        guard_str = " ∧ ".join([str(atom) for atom in self.guard]).strip()
        keys_str = ", ".join([str(var) for var in self.keys]).strip()
        values_str = ", ".join([str(var) for var in self.values]).strip()
        return f"{guard_str}: {{{keys_str}}} -> {{{values_str}}}"

    # computes a canonical string representation of the functional dependency
    def _compute_canonical_str(self) -> str:
        # we construct a stable tuple-term representation of the fd
        fd_term = ("fd",)
        for atom in sorted(self.guard, key=lambda x: x.relation.name):
            atom_term = (atom.relation.name,)
            for arg in atom.args:
                if isinstance(arg, Var):
                    atom_term += (("var", arg.name, str(arg.type)),)
                elif isinstance(arg, Literal):
                    atom_term += (("lit", arg.value, str(arg.type)),)
                else:
                    atom_term += (("arg", str(arg)),)
            fd_term += (atom_term,)
        keys_term = tuple(sorted((("var", v.name, str(v.type)) for v in self.keys)))
        values_term = tuple(sorted((("var", v.name, str(v.type)) for v in self.values)))
        fd_term += (("keys", keys_term), ("values", values_term))
        return str(fd_term)

    @property
    def canonical_str(self) -> str:
        """
        A canonical string representation (depends on guard atoms, keys, and values).
        """
        return self._canonical_str

#
# A unary guard atom `T(x::T)` is redundant if the constraint contains a non-unary atom
# `R(...,x::T,...)`.  We discard all redundant guard atoms in the constructed fd.
#
# TODO: to gracefully handle Require nodes with non and multiple FDs, change this method to
# normalized_fds(require: Require) -> Sequence[FunctionalDependency]
def normalized_fd(lookup: Lookup) -> Optional[FunctionalDependency]:
    """
    If the input `Lookup` is a valid `unique_fields` constraint, constructs its reduced
    functional dependency `φ: X -> Y`, where `φ` contains all non-redundant guard atoms,
    `X` are the variables used in the unique fields, and `Y` are the remaining variables.
    Returns `None` if the input lookup is not a valid `unique_fields` constraint.
    """
    parts = _split_unique_fields_lookup(lookup)
    if parts is None:
        return None
    all_vars, unique_vars, guard_atoms = parts

    # remove redundant lookups
    redundant_guard_atoms: list[Lookup] = []
    for atom in guard_atoms:
        # the atom is unary A(x::T)
        if len(atom.args) != 1:
            continue
        var = atom.args[0]
        assert isinstance(var, Var)
        # T is a scalar type (which includes entity types)
        var_type = var.type
        if not isinstance(var_type, ScalarType):
            continue
        # the atom is a entity typing T(x::T) i.e., T = A (and hence not a Boolean property)
        var_type_name = var_type.name
        rel_name = atom.relation.name
        if rel_name != var_type_name:
            continue
        # Found an atom of the form T(x::T)
        # check if there is another atom R(...,x::T,...)
        for typed_atom in guard_atoms:
            if len(typed_atom.args) == 1:
                continue
            if var in typed_atom.args:
                redundant_guard_atoms.append(atom)
                break

    guard = [atom for atom in guard_atoms if atom not in redundant_guard_atoms]
    keys = unique_vars
    values = [v for v in all_vars if v not in keys]

    return FunctionalDependency(guard, keys, values)

# TODO: this will need to be migrated later when needed for later passes, and may also need
# re-working.
#
# def contains_only_declarable_constraints(node: Node) -> bool:
#     """
#     Checks whether the input node contains only `Require` nodes annotated with
#     `declare_constraint` (or such a node itself).
#     """
#     # Check if the node itself is a Require node with declarable constraint
#     if isinstance(node, Require):
#         return is_declarable_constraint(node)
#
#     # Otherwise, check if it is a Logical node containing only declarable constraints
#     if not isinstance(node, Logical):
#         return False
#     if len(node.body) == 0:
#         return False
#     for task in node.body:
#         if not contains_only_declarable_constraints(task):
#             return False
#     return True
#
# def is_declarable_constraint(node: Node) -> bool:
#     """
#     Checks whether the input node is a `Require` node annotated with `declare_constraint`.
#     """
#     if not isinstance(node, Require):
#         return False
#     return builtins.declare_constraint_annotation in node.annotations
