from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
from ....metamodel.metamodel import Require, Logical, Lookup, Task
from ....metamodel.rewriter import Rewriter
from ....metamodel.builtins import builtins
from .functional_dependencies import normalized_fd
from .lqp_pass import Pass
from .....config import Config
from .. import annotations as lqp_annotations


class AnnotateConstraints(Pass):
    """
    Extends discharge constraints logic by annotating Require nodes based on how they
    should be treated when generating code:
     * `@declare_constraint` if the Require represents a constraint that can be declared in LQP
     * `@discharged` if the Require represents a constraint that should be dismissed during
       code generation. Namely, when it cannot be declared in LQP and uses one of the
       `unique`, `exclusive`, `anyof` builtins. These nodes are removed from the IR in
       later passes.

    By default, all constraints are discharged. To enable emitting constraints, set the
    `reasoners.logic.emit_constraints` flag to True in the config file:
    ```yaml
    reasoners:
      logic:
        emit_constraints: true
    ```
    """

    def _rewrite_task(self, node: Task, config: Config) -> Task:
        rewriter = AnnotateConstraintsRewriter(emit_constraints=config.reasoners.logic.emit_constraints)
        return rewriter(node)

@dataclass
class AnnotateConstraintsRewriter(Rewriter):
    """
    Visitor marks Require nodes with:
    - `@discharged` if they should be discharged from the metamodel
    - `@declare_constraint` if they should be kept and emitted as LQP constraint declarations

    By default, all constraints are discharged. To enable emitting constraints, set the
    `emit_constraints` flag to True.
    """
    def __init__(self, emit_constraints: bool = False):
        super().__init__()
        self.emit_constraints = emit_constraints

    def require(self, node: Require) -> Optional[Require]:
        """
        Annotate Require nodes with either @declare_constraint or @discharged.

        First checks if the constraint should be declared (if emit_constraints is enabled).
        Otherwise, discharges constraints that use unique, exclusive, or anyof builtins.
        """
        # First, check if we should declare this constraint
        if self._should_declare_constraint(node):
            # Check if already annotated
            if any(a.relation.name == lqp_annotations.declare_constraint.relation.name for a in node.annotations):
                return None  # Already annotated

            # Add declare_constraint annotation
            new_annotations = node.annotations + (lqp_annotations.declare_constraint,)
            return node.mut(annotations=new_annotations)

        # Otherwise, check if we should discharge it
        if self._should_discharge(node):
            # Check if already annotated
            if any(a.relation.name == lqp_annotations.discharged.relation.name for a in node.annotations):
                return None  # Already annotated

            # Add discharge annotation
            new_annotations = node.annotations + (lqp_annotations.discharged,)
            return node.mut(annotations=new_annotations)

        return None  # No change

    def _should_declare_constraint(self, node: Require) -> bool:
        """
        Check if a Require node should be declared as a constraint in LQP.

        A constraint should be declared if:
        1. emit_constraints is enabled
        2. It's a valid unique constraint (has domain with guards and check with unique)
        3. The functional dependency is non-structural (structural ones are handled by @function)
        """
        if not self.emit_constraints:
            return False

        # Only declare constraints for unique constraints
        if not self._is_valid_unique_constraint(node):
            return False

        # Currently, we only declare non-structural functional dependencies
        fd = normalized_fd(node)
        return fd is not None and not fd.is_structural

    def _should_discharge(self, node: Require) -> bool:
        """
        Check if a Require node should be discharged.

        A constraint should be discharged if the check contains one of the constraint
        builtins: unique, exclusive, or anyof.
        """
        # Check if the check contains dischargeable constraint builtins
        check = node.check
        if not isinstance(check, Logical):
            return False

        discharged_names = [
            builtins.constraints.unique.name,
            builtins.constraints.exclusive.name,
            builtins.constraints.anyof.name
        ]

        for task in check.body:
            if isinstance(task, Lookup) and task.relation.name in discharged_names:
                return True

        return False

    def _is_valid_unique_constraint(self, node: Require) -> bool:
        """
        Check if a Require node is a valid unique constraint.

        A valid unique constraint has:
        - A domain (can be empty, Logical with lookups, or a single Lookup)
        - A check that is a Logical containing a unique lookup
        """
        # Check must be a Logical node
        if not isinstance(node.check, Logical):
            return False

        # Look for a unique lookup in the check
        has_unique = False
        for task in node.check.body:
            if isinstance(task, Lookup) and task.relation.name == builtins.constraints.unique.name:
                has_unique = True
                break

        return has_unique
