# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: F401

from typing import Optional

import pytest
from tvm_ffi.access_path import AccessPath

from tvm.script import tirx as T


@T.prim_func(s_tir=True)
def _func():
    T.evaluate(-1)
    T.evaluate(1)
    T.evaluate(2)
    T.evaluate(3)
    T.evaluate(4)
    T.evaluate(5)
    T.evaluate(6)
    T.evaluate(7)


def test_annotation_multi_access_paths():
    result = _func.with_attr("global_symbol", "main").script(
        path_to_annotate={
            AccessPath.root().attr("body").attr("seq").array_item(1): "annotation 1",
            AccessPath.root().attr("body").attr("seq").array_item(3): "annotation 3",
            AccessPath.root().attr("body").attr("seq").array_item(5): "annotation 5",
            AccessPath.root().attr("body").attr("seq").array_item(7): "annotation 7",
        }
    )
    assert (
        result
        == """# from tvm.script import tirx as T
# from tvm.tirx.layout import Axis

@T.prim_func(s_tir=True)
def main():
    T.evaluate(-1)
    T.evaluate(1)  # annotation 1
    T.evaluate(2)
    T.evaluate(3)  # annotation 3
    T.evaluate(4)
    T.evaluate(5)  # annotation 5
    T.evaluate(6)
    T.evaluate(7)  # annotation 7"""
    )


def test_annotate_from_multi_obj():
    result = _func.with_attr("global_symbol", "main").script(
        obj_to_annotate={
            _func.body.seq[1]: "annotation 1",
            _func.body.seq[3]: "annotation 3",
            _func.body.seq[5]: "annotation 5",
            _func.body.seq[7]: "annotation 7",
        }
    )
    assert (
        result
        == """# from tvm.script import tirx as T
# from tvm.tirx.layout import Axis

@T.prim_func(s_tir=True)
def main():
    T.evaluate(-1)
    T.evaluate(1)  # annotation 1
    T.evaluate(2)
    T.evaluate(3)  # annotation 3
    T.evaluate(4)
    T.evaluate(5)  # annotation 5
    T.evaluate(6)
    T.evaluate(7)  # annotation 7"""
    )


def test_disable_concise_scoping_when_scope_annotated():
    @T.prim_func(s_tir=True)
    def _func():
        x = 1
        y = x + 1
        T.evaluate(y - 1)

    # In fork, each bare `x = expr` lowers to AllocBuffer + BufferStore (local_scalar);
    # the printer fuses each pair into a single `y: T.int32 = x + 1` line. Annotate the
    # AllocBuffer that originates this fused line.
    result = _func.with_attr("global_symbol", "main").script(
        obj_to_annotate={
            _func.body.seq[2]: "annotation 1",
        }
    )
    assert (
        result
        == """# from tvm.script import tirx as T
# from tvm.tirx.layout import Axis

@T.prim_func(s_tir=True)
def main():
    x: T.int32 = 1
    y: T.int32 = x + 1  # annotation 1
    T.evaluate(y - 1)"""
    )