from __future__ import annotations
from typing import Optional, Sequence
from ....metamodel.metamodel import (
    Require, Logical, Var, Relation, Lookup, ScalarType, Literal
)
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` contains one or more relation `Lookup` tasks (the guard atoms)
 * the `Require` node's `check` is a `Logical` node containing a single `unique` lookup
 * the `unique` lookup has precisely one argument, which is a tuple of `Var` objects
 * the variables in the `unique` tuple must be a subset of the variables in the domain

Note: The TransformRequire rewrite pass (which runs first in the pipeline) normalizes all
unique constraints to this format. By the time this code runs, all Require nodes with unique
constraints are guaranteed to be in this standard format.
============================================================================================

We use the following unique constraint as the running example (standard format):

```
Require
    domain
        Logical
            Person(person::Person)
            first_name(person::Person, firstname::String)
            last_name(person::Person, lastname::String)
    check
        Logical
            unique((firstname::String, lastname::String))
```

This represents a schema with a unary `Person` entity relation and two binary relations
`first_name` and `last_name` that relate persons to their names.

=========================== 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 predicates in the domain.
============================================================================================

In the running example, the domain 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 the
combination of first and last name must be unique (i.e., determines the person).

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 typing atoms during
compilation. For instance, in the running example, the atom `Person(person::Person)` in the
domain is redundant because the type of the `person` variable is already specified in the
other two atoms `first_name` and `last_name`. Consequently, we identify and remove redundant
atoms from the functional dependency representation.

Formally, a _guard_ atom is any `Lookup` node in the domain. A unary guard atom `T(x::T)` is
_redundant_ if the domain also contains a non-unary guard atom `R(...,x::T,...)` with the
same variable.

================================ Normalized FDs ============================================
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 from the domain
 2. `X` is the set of variables used in the `unique` atom
 3. `Y` is the set of all other variables used in the guard atoms
============================================================================================

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}`.
"""

def normalized_fd(node: Require) -> Optional[FunctionalDependency]:
    """
    Constructs a functional dependency from a `Require` node in the new format where:
    - The domain contains the guard atoms (relation lookups)
    - The check contains a single `unique((var1, var2, ...))` call

    Returns `None` if the node doesn't match this format.
    """
    # Extract guard atoms from domain (can be a Logical node or a single Lookup)
    guard_atoms: list[Lookup] = []
    if isinstance(node.domain, Logical):
        for task in node.domain.body:
            if isinstance(task, Lookup):
                guard_atoms.append(task)
            else:
                # Domain should only contain lookups
                return None
    elif isinstance(node.domain, Lookup):
        guard_atoms.append(node.domain)
    else:
        return None

    if len(guard_atoms) == 0:
        return None

    # Find the unique lookup in check (can be a Logical node or a single Lookup)
    unique_lookup: Optional[Lookup] = None
    if isinstance(node.check, Logical):
        for task in node.check.body:
            if isinstance(task, Lookup) and task.relation.name == builtins.constraints.unique.name:
                if unique_lookup is not None:
                    # Multiple unique calls not supported
                    return None
                unique_lookup = task
    elif isinstance(node.check, Lookup):
        if node.check.relation.name == builtins.constraints.unique.name:
            unique_lookup = node.check
    else:
        return None

    if unique_lookup is None:
        return None

    # Extract unique variables from unique((var1, var2, ...))
    if len(unique_lookup.args) != 1:
        return None

    unique_arg = unique_lookup.args[0]
    if not isinstance(unique_arg, tuple):
        return None

    unique_vars: list[Var] = []
    for arg in unique_arg:
        if not isinstance(arg, Var):
            return None
        unique_vars.append(arg)

    if len(unique_vars) == 0:
        return None

    # Collect all variables from guard atoms
    all_vars: list[Var] = []
    for atom in guard_atoms:
        for arg in atom.args:
            if isinstance(arg, Var) and arg not in all_vars:
                all_vars.append(arg)

    # Remove redundant guard atoms (unary typing atoms T(x::T) where x appears elsewhere)
    redundant_guard_atoms: list[Lookup] = []
    for atom in guard_atoms:
        # Check if atom is unary A(x::T)
        if len(atom.args) != 1:
            continue
        var = atom.args[0]
        if not isinstance(var, Var):
            continue
        # Check if T is a scalar type
        var_type = var.type
        if not isinstance(var_type, ScalarType):
            continue
        # Check if atom is entity typing T(x::T) (i.e., T = A)
        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)

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

# 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
