import os
from collections.abc import Callable
from dataclasses import dataclass

from xdsl.context import Context
from xdsl.dialects import builtin, pdl_interp
from xdsl.interpreter import (
    Interpreter,
    InterpreterFunctions,
    PythonValues,
    impl_external,
    register_impls,
)
from xdsl.interpreters.eqsat_pdl_interp import EqsatPDLInterpFunctions
from xdsl.interpreters.pdl_interp import PDLInterpFunctions
from xdsl.ir import Operation
from xdsl.parser import Parser
from xdsl.passes import ModulePass
from xdsl.pattern_rewriter import PatternRewriter
from xdsl.traits import SymbolTable

_DEFAULT_MAX_ITERATIONS = 20
"""Default number of times to iterate over the module."""


# TODO: remove the constraint functions here (https://github.com/xdslproject/xdsl/issues/5391)
@register_impls
class EqsatConstraintFunctions(InterpreterFunctions):
    @impl_external("is_not_unsound")
    def run_is_not_unsound(
        self, interp: Interpreter, _op: Operation, args: PythonValues
    ):
        assert isinstance(op := args[0], Operation)
        return "unsound" not in op.attributes, ()


def apply_eqsat_pdl_interp(
    op: builtin.ModuleOp,
    ctx: Context,
    pdl_interp_module: builtin.ModuleOp,
    max_iterations: int = _DEFAULT_MAX_ITERATIONS,
    callback: Callable[[builtin.ModuleOp], None] | None = None,
):
    matcher = SymbolTable.lookup_symbol(pdl_interp_module, "matcher")
    assert isinstance(matcher, pdl_interp.FuncOp)
    assert matcher is not None, "matcher function not found"

    # Initialize interpreter and implementations once
    interpreter = Interpreter(pdl_interp_module)
    pdl_interp_functions = PDLInterpFunctions()
    eqsat_pdl_interp_functions = EqsatPDLInterpFunctions()
    PDLInterpFunctions.set_ctx(interpreter, ctx)
    eqsat_pdl_interp_functions.populate_known_ops(op)
    interpreter.register_implementations(eqsat_pdl_interp_functions)
    interpreter.register_implementations(pdl_interp_functions)
    interpreter.register_implementations(EqsatConstraintFunctions())

    if not op.ops.first:
        return

    rewriter = PatternRewriter(op.ops.first)
    rewriter.operation_modification_handler.append(
        eqsat_pdl_interp_functions.modification_handler
    )
    pdl_interp_functions.set_rewriter(interpreter, rewriter)
    for _i in range(max_iterations):
        for root in op.body.walk():
            rewriter.current_operation = root
            interpreter.call_op(matcher, (root,))
        eqsat_pdl_interp_functions.execute_pending_rewrites(interpreter)

        if not eqsat_pdl_interp_functions.worklist:
            break

        eqsat_pdl_interp_functions.rebuild(interpreter)
        if callback is not None:
            callback(op)


@dataclass(frozen=True)
class ApplyEqsatPDLInterpPass(ModulePass):
    name = "apply-eqsat-pdl-interp"

    pdl_interp_file: str | None = None
    max_iterations: int = _DEFAULT_MAX_ITERATIONS
    """Maximum number of iterations to run, default 20."""

    def apply(self, ctx: Context, op: builtin.ModuleOp) -> None:
        if self.pdl_interp_file is not None:
            assert os.path.exists(self.pdl_interp_file)
            with open(self.pdl_interp_file) as f:
                pdl_interp_module_str = f.read()
                parser = Parser(ctx, pdl_interp_module_str)
                pdl_interp_module = parser.parse_module()
        else:
            pdl_interp_module = op

        apply_eqsat_pdl_interp(op, ctx, pdl_interp_module, self.max_iterations)
