from __future__ import annotations
from dataclasses import dataclass, field
from relationalai.clients.config import Config
from relationalai.semantics.metamodel import builtins
from relationalai.semantics.metamodel.ir import Node, Model, Require
from relationalai.semantics.metamodel.compiler import Pass
from relationalai.semantics.metamodel.rewrite.discharge_constraints import (
    DischargeConstraintsVisitor
)
from relationalai.semantics.lqp.rewrite.functional_dependencies import (
    normalized_fds
)

class AnnotateConstraints(Pass):
    """
    Extends `DischargeConstraints` pass by discharging only those Require nodes that cannot
    be declared as constraints in LQP.

    More precisely, the pass annotates Require nodes depending on how they should be
    treated when generating code:
     * `@declare_constraint` if the Require represents a constraint that can be declared in LQP.
     * `@discharge` 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 model
       in the Flatten pass.
    """

    def rewrite(self, model: Model, options: dict = {}) -> Model:
        # Extract config from options if provided (for compatibility with PyRel ShimExecutor)
        config = options.get('_v0_config', None)
        return AnnotateConstraintsRewriter(config=config).walk(model)

@dataclass
class AnnotateConstraintsRewriter(DischargeConstraintsVisitor):
    emit_constraints: bool = field(default=False)
    config: Config | None = field(default=None, repr=False)
    """
    Visitor marks Require nodes with
    - `discharge` 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
    `reasoner.rule.emit_constraints` flag to True in the config file.
    ```toml
    [reasoner.rule]
    emit_constraints = true
    ```
    """

    def __post_init__(self):
        from relationalai.semantics.internal.internal import overridable_flag
        # Use provided config if available, otherwise try to create one with fetch=False
        # to avoid requiring raiconfig.toml when config is passed programmatically
        if self.config is None:
            try:
                self.config = Config(fetch=False)
            except Exception:
                # If Config() fails even with fetch=False, use empty dict
                self.config = Config({}, fetch=False)
        self.emit_constraints = overridable_flag('reasoner.rule.emit_constraints', self.config, None, False)

    def _should_declare_constraint(self, node: Require) -> bool:
        if not self.emit_constraints:
            return False
        fds = normalized_fds(node)
        if len(fds) == 0:
            return False
        # Currently, we only declare non-structural functional dependencies.
        return any(not fd.is_structural for fd in fds)

    def handle_require(self, node: Require, parent: Node):
        if self._should_declare_constraint(node):
            return node.reconstruct(
                node.engine,
                node.domain,
                node.checks,
                node.annotations | [builtins.declare_constraint_annotation]
            )

        return super().handle_require(node, parent)
