from .....config import Config
from ....metamodel.metamodel import Require, Logical, Not, Task

from ..rewriter import RewriterWithIntermediates
from ..annotations import discharged, declare_constraint
from .lqp_pass import Pass


class HandleAnnotatedRequires(Pass):
    """
    Processes Require nodes based on their annotations, similar to v0 Flatten.handle_require.

    This pass handles three types of Require nodes:
    1. `@discharged ` annotated: Removed completely
    2. `@declare_constraint` annotated: Kept as-is (these will be declared as LQP constraints)
    3. Other requires with errors: Generate error-checking logic extracted to top-level
    """

    def _rewrite_task(self, node: Task, config: Config) -> Task:
        rewriter = HandleRequireRewriter(node)
        result = rewriter(node)
        return result


class HandleRequireRewriter(RewriterWithIntermediates):
    """
    Rewriter that processes Require nodes and extracts error logic to top level.
    """

    def __init__(self, node: Task):
        assert isinstance(node, Logical), f"Expected a root logical, but was given {type(node)}"
        super().__init__(node)

    def logical(self, node: Logical):
        # We need to rewrite at the Logical level to be able to remove Require nodes that
        # are discharged
        new_body = []
        changed = False

        for task in node.body:
            if isinstance(task, Require):
                changed = True
                handled = self._handle_require(task)
                if handled is not None:
                    new_body.append(handled)
            elif isinstance(task, Logical) and not task.body:
                changed = True
                # Skip empty logicals. These can come about because of the structure where
                # `Require` nodes are nested inside of a `Logical`, and they may be dropped
                # leaving an empty `Logical`.
                continue
            else:
                new_body.append(task)

        if changed:
            result = node.mut(body=tuple(new_body))
        else:
            result = node

        return super().logical(result)

    def _handle_require(self, node: Require):
        """
        Handle a single Require node based on its annotations.

        Similar to v0 Flatten.handle_require:
        - If @discharged: remove (return None)
        - If @declare_constraint: keep as-is (return node)
        - Otherwise: generate error logic to top_level and remove (return None)
        """
        # Check if this Require is discharged
        if any(a.relation.name == discharged.relation.name for a in node.annotations):
            # Remove discharged Requires
            return None

        # Check if this Require is a declare_constraint
        if any(a.relation.name == declare_constraint.relation.name for a in node.annotations):
            # Keep Requires that are declared constraints. They are handled later when
            # generating LQP
            return node

        # Generate error logic for remaining requires (if they have errors)
        if node.error:
            # Generate: domain ∧ ¬check ⇒ error
            error_tasks = []

            # Add domain tasks
            error_tasks.append(node.domain)

            # Add negated check
            error_tasks.append(Not(task=node.check, reasoner=node.reasoner, source=node.source))

            # Add error tasks
            error_tasks.append(node.error)

            # Create top-level rule
            self.add_intermediate_task(
                Logical(body=tuple(error_tasks), reasoner=node.reasoner, source=node.source, scope=True)
            )

        return None
