# RUN: %PYTHON %s | FileCheck %s
# REQUIRES: python-ge-311
# UNSUPPORTED: python-stable-abi
import gc
from contextlib import contextmanager

from mlir.ir import *
from mlir.dialects._ods_common import _cext
from mlir.dialects import arith, _arith_ops_gen


def run(f):
    print("\nTEST:", f.__name__)
    f()
    gc.collect()
    assert Context._get_live_count() == 0


# CHECK-LABEL: TEST: testInferLocations
@run
def testInferLocations():
    with Context() as ctx, loc_tracebacks():
        ctx.allow_unregistered_dialects = True

        op = Operation.create("custom.op1")
        one = arith.constant(IndexType.get(), 1)
        _cext.globals.register_traceback_file_exclusion(arith.__file__)
        two = arith.constant(IndexType.get(), 2)

        # fmt: off
        # CHECK: loc(callsite("testInferLocations"("{{.*}}[[SEP:[/\\]+]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:13 to :43) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":14:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:1 to :4))))
        # fmt: on
        print(op.location)

        # Test nesting of loc_tracebacks().
        with loc_tracebacks():
            # fmt: off
            # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]arith.py":45:12 to :76) at callsite("constant"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]arith.py":90:40 to :81) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:14 to :48) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:1 to :4))))))
            # fmt: on
            print(one.location)

        # fmt: off
        # CHECK: loc(callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:14 to :48) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:1 to :4))))
        # fmt: on
        print(two.location)

        _cext.globals.register_traceback_file_inclusion(_arith_ops_gen.__file__)
        three = arith.constant(IndexType.get(), 3)
        # fmt: off
        # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":{{[0-9]+}}:4 to :235) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:16 to :50) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:1 to :4)))))
        # fmt: on
        print(three.location)

        def foo():
            four = arith.constant(IndexType.get(), 4)
            print(four.location)

        # fmt: off
        # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":{{[0-9]+}}:4 to :235) at callsite("testInferLocations.<locals>.foo"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:19 to :53) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:8 to :13) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:1 to :4))))))
        # fmt: on
        foo()

        _cext.globals.register_traceback_file_exclusion(__file__)

        # fmt: off
        # CHECK: loc("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":{{[0-9]+}}:4 to :235))
        # fmt: on
        foo()

        def bar1():
            def bar2():
                def bar3():
                    five = arith.constant(IndexType.get(), 5)
                    print(five.location)

                bar3()

            bar2()

        _cext.globals.register_traceback_file_inclusion(__file__)
        _cext.globals.register_traceback_file_exclusion(_arith_ops_gen.__file__)

        _cext.globals.set_loc_tracebacks_frame_limit(2)
        # fmt: off
        # CHECK: loc(callsite("testInferLocations.<locals>.bar1.<locals>.bar2.<locals>.bar3"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:27 to :61) at "testInferLocations.<locals>.bar1.<locals>.bar2"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:16 to :22)))
        # fmt: on
        bar1()

        _cext.globals.set_loc_tracebacks_frame_limit(1)
        # fmt: off
        # CHECK: loc("testInferLocations.<locals>.bar1.<locals>.bar2.<locals>.bar3"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":{{[0-9]+}}:27 to :61))
        # fmt: on
        bar1()

        _cext.globals.set_loc_tracebacks_frame_limit(0)
        # CHECK: loc(unknown)
        bar1()


# CHECK-LABEL: TEST: testNamelocWrap
@run
def testNamelocWrap():
    with Context() as ctx, Location.unknown():
        ctx.allow_unregistered_dialects = True
        with loc_tracebacks(current_loc_actn=CurrentLocAction.NAMELOC_WRAP):
            # Build nested NameLoc: TaskA(ResourceB(unknown))
            inner = Location.name("ResourceB", childLoc=Location.unknown())
            outer = Location.name("TaskA", childLoc=inner)
            with outer:
                op = Operation.create("custom.op1")
                # fmt: off
                # CHECK: loc("TaskA"("ResourceB"(callsite(
                # fmt: on
                print(op.location)


# CHECK-LABEL: TEST: testOnExplicitDefault
@run
def testOnExplicitDefault():
    with Context() as ctx, Location.unknown():
        ctx.allow_unregistered_dialects = True
        with loc_tracebacks():
            explicit = Location.file("explicit.py", 1, 1)
            op = Operation.create("custom.op1", loc=explicit)
            # CHECK: loc("explicit.py":1:1)
            print(op.location)


# CHECK-LABEL: TEST: testOnExplicitUseTraceback
@run
def testOnExplicitUseTraceback():
    with Context() as ctx, Location.unknown():
        ctx.allow_unregistered_dialects = True
        with loc_tracebacks(on_explicit_actn=OnExplicitAction.USE_TRACEBACK):
            explicit = Location.file("explicit.py", 1, 1)
            op = Operation.create("custom.op1", loc=explicit)
            # fmt: off
            # CHECK: loc(callsite(
            # fmt: on
            print(op.location)


# CHECK-LABEL: TEST: testDslProfilingUseCase
@run
def testDslProfilingUseCase():
    with Context() as ctx, Location.unknown():
        ctx.allow_unregistered_dialects = True
        with loc_tracebacks(
            current_loc_actn=CurrentLocAction.NAMELOC_WRAP,
            on_explicit_actn=OnExplicitAction.USE_TRACEBACK,
        ):
            task_loc = Location.name("Task", childLoc=Location.unknown())
            with task_loc:
                op1 = Operation.create("custom.op1", loc=Location.file("x.py", 1, 1))
                # fmt: off
                # CHECK: loc("Task"(callsite(
                # fmt: on
                print(op1.location)

                op2 = Operation.create("custom.op2")
                # fmt: off
                # CHECK: loc("Task"(callsite(
                # fmt: on
                print(op2.location)


# CHECK-LABEL: TEST: testOnExplicitFallbackWhenTracebacksDisabled
@run
def testOnExplicitFallbackWhenTracebacksDisabled():
    """When on_explicit != use_explicit but tracebacks are disabled,
    the explicit loc should be returned (not Location.current)."""
    with Context() as ctx, Location.unknown():
        ctx.allow_unregistered_dialects = True
        # Set on_explicit to USE_TRACEBACK WITHOUT enabling tracebacks.
        _cext.globals.set_traceback_action_on_explicit_loc(
            OnExplicitAction.USE_TRACEBACK
        )
        try:
            explicit = Location.file("keep_me.py", 42, 1)
            op = Operation.create("custom.op1", loc=explicit)
            # CHECK: loc("keep_me.py":42:1)
            print(op.location)
        finally:
            _cext.globals.set_traceback_action_on_explicit_loc(
                OnExplicitAction.USE_EXPLICIT
            )


# CHECK-LABEL: TEST: testUseExplicitWithNamelocWrap
@run
def testUseExplicitWithNamelocWrap():
    """on_explicit=use_explicit + current_loc=nameloc_wrap should wrap
    the explicit loc with the NameLoc chain (flags are orthogonal)."""
    with Context() as ctx, Location.unknown():
        ctx.allow_unregistered_dialects = True
        with loc_tracebacks(
            on_explicit_actn=OnExplicitAction.USE_EXPLICIT,
            current_loc_actn=CurrentLocAction.NAMELOC_WRAP,
        ):
            task_loc = Location.name("Task", childLoc=Location.unknown())
            with task_loc:
                explicit = Location.file("framework.py", 10, 1)
                op = Operation.create("custom.op1", loc=explicit)
                # CHECK: loc("Task"("framework.py":10:1))
                print(op.location)