from typing import Optional
from ....metamodel.metamodel import (
    Require, Relation, Annotation, Task, Update
)
from ....metamodel.rewriter import Rewriter, Walker
from .functional_dependencies import normalized_fd
from .lqp_pass import Pass
from .....config import Config
from .. import annotations as lqp_annotations

class FunctionAnnotations(Pass):
    """
    Pass marks all appropriate relations with `function` annotation. Collects functional
    dependencies from unique Require nodes and uses this information to identify functional
    relations.
    """

    def _rewrite_task(self, node: Task, config: Config) -> Task:
        # Collect functional dependencies from all unique constraints
        collect_fds = CollectFDsVisitor()
        collect_fds(node)

        # Apply function annotations to relations and updates
        rewriter = FunctionalAnnotationsRewriter(collect_fds.functional_relations)
        annotated_task = rewriter(node)
        return annotated_task


class CollectFDsVisitor(Walker):
    """
    Visitor collects all unique constraints.
    """

    # Currently, only information about k-functional fd is collected.
    def __init__(self):
        super().__init__()
        self.functional_relations:dict[Relation, int] = {}

    def require(self, node: Require):
        # Try the new format first (domain with guard lookups, check with unique)
        fd = normalized_fd(node)
        if fd is not None and fd.is_structural:
            relation = fd.structural_relation
            k = fd.structural_rank
            current_k = self.functional_relations.get(relation, 0)
            self.functional_relations[relation] = max(current_k, k)
            return

class FunctionalAnnotationsRewriter(Rewriter):
    """
    This visitor marks functional_relations with `@function(:checked [, k])` annotation.
    """

    def __init__(self, functional_relations: dict[Relation, int]):
        super().__init__()
        self.functional_relations = functional_relations

    def get_functional_annotation(self, rel: Relation) -> Optional[Annotation]:
        k = self.functional_relations.get(rel, None)
        if k is None:
            return None
        if k == 1:
            return lqp_annotations.function_checked
        return lqp_annotations.function_ranked_checked(k)

    def relation(self, node: Relation) -> Optional[Relation]:
        """Add function annotation to relations with functional dependencies."""
        function_annotation = self.get_functional_annotation(node)
        if function_annotation is None:
            return None  # No change

        # Check if annotation already exists
        if any(a.relation.name == function_annotation.relation.name for a in node.annotations):
            return None  # Already annotated

        # Add the function annotation
        new_annotations = node.annotations + (function_annotation,)
        return node.mut(annotations=new_annotations)

    def update(self, node: Update) -> Optional[Update]:
        """Add function annotation to updates on functional relations."""
        function_annotation = self.get_functional_annotation(node.relation)
        if function_annotation is None:
            return None  # No change

        # Check if annotation already exists
        if any(a.relation.name == function_annotation.relation.name for a in node.annotations):
            return None  # Already annotated

        # Add the function annotation
        new_annotations = node.annotations + (function_annotation,)
        return node.mut(annotations=new_annotations)
