"""
Test module for proton's Python API.
No GPU kernel should be declared in this test.
Profile correctness tests involving GPU kernels should be placed in `test_profile.py`.
"""

import pytest
import json
import torch
import triton.profiler as proton
import pathlib
from triton.profiler.hooks.hook import HookManager
from triton.profiler.hooks.launch import LaunchHook
from triton.profiler.hooks.instrumentation import InstrumentationHook
from triton.profiler.metric import transform_tensor_metrics
from triton._internal_testing import is_hip


def test_profile_single_session(tmp_path: pathlib.Path):
    temp_file0 = tmp_path / "test_profile0.hatchet"
    session_id0 = proton.start(str(temp_file0.with_suffix("")))
    proton.activate()
    proton.deactivate()
    proton.finalize()
    assert temp_file0.exists()

    temp_file1 = tmp_path / "test_profile1.hatchet"
    session_id1 = proton.start(str(temp_file1.with_suffix("")))
    proton.activate(session_id1)
    proton.deactivate(session_id1)
    proton.finalize(session_id1)
    assert session_id1 == session_id0 + 1
    assert temp_file1.exists()

    session_id2 = proton.start("test")
    proton.activate(session_id2)
    proton.deactivate(session_id2)
    proton.finalize()
    assert session_id2 == session_id1 + 1
    assert pathlib.Path("test.hatchet").exists()
    pathlib.Path("test.hatchet").unlink()


def test_roofline_scalar_metric_types():
    scalar_metrics, tensor_metrics = transform_tensor_metrics({
        "flops": 1,
        "flops16": 2,
        "flops32": [5, 6],
        "flops64": torch.tensor([7, 8], dtype=torch.int64),
        "bytes": 3.0,
        "custom": 4,
    })

    assert tensor_metrics == {}
    assert scalar_metrics["flops"] == 1.0
    assert type(scalar_metrics["flops"]) is float
    assert scalar_metrics["flops16"] == 2.0
    assert type(scalar_metrics["flops16"]) is float
    assert scalar_metrics["flops32"] == [5.0, 6.0]
    assert scalar_metrics["flops64"] == [7.0, 8.0]
    assert scalar_metrics["bytes"] == 3
    assert type(scalar_metrics["bytes"]) is int
    assert scalar_metrics["custom"] == 4
    assert type(scalar_metrics["custom"]) is int

    scalar_metrics, tensor_metrics = transform_tensor_metrics({"bytes": torch.tensor([9.0, 10.0])})
    assert tensor_metrics == {}
    assert scalar_metrics["bytes"] == [9, 10]


def test_profile_multiple_sessions(tmp_path: pathlib.Path):
    temp_file0 = tmp_path / "test_profile0.hatchet"
    proton.start(str(temp_file0.with_suffix("")))
    temp_file1 = tmp_path / "test_profile1.hatchet"
    proton.start(str(temp_file1.with_suffix("")))
    proton.activate()
    proton.deactivate()
    proton.finalize()
    assert temp_file0.exists()
    assert temp_file1.exists()

    temp_file2 = tmp_path / "test_profile2.hatchet"
    session_id2 = proton.start(str(temp_file2.with_suffix("")))
    temp_file3 = tmp_path / "test_profile3.hatchet"
    session_id3 = proton.start(str(temp_file3.with_suffix("")))
    proton.deactivate(session_id2)
    proton.deactivate(session_id3)
    proton.finalize()
    assert temp_file2.exists()
    assert temp_file3.exists()


def test_profile_mode(tmp_path: pathlib.Path):
    temp_file0 = tmp_path / "test_profile0.hatchet"
    if is_hip():
        try:
            proton.start(str(temp_file0.with_suffix("")), mode="pcsampling")
        except Exception as e:
            assert "unsupported mode: pcsampling" in str(e)
        finally:
            proton.finalize()
    else:
        import os
        import pytest

        if os.environ.get("PROTON_SKIP_PC_SAMPLING_TEST", "0") == "1":
            pytest.skip("PC sampling test is disabled")

        # Two sessions with the same mode can coexist
        proton.start(str(temp_file0.with_suffix("")), mode="pcsampling")
        temp_file1 = tmp_path / "test_profile1.hatchet"
        proton.start(str(temp_file1.with_suffix("")), mode="pcsampling")
        proton.finalize()
        assert temp_file1.exists()

        # Two sessions with different modes cannot coexist
        try:
            proton.start(str(temp_file0.with_suffix("")), mode="pcsampling")
            proton.start(str(temp_file1.with_suffix("")))
        except Exception as e:
            assert "Cannot add a session with the same profiler but a different mode than existing sessions" in str(e)
        finally:
            proton.finalize()

        # Two sessions with different modes cannot coexist even if the first session is deactivated.
        # In proton, once we deactivate a session, its profiler is not stopped, so changing the profiler mode is not allowed
        # The only way to start a session with a different mode is to finalize all existing sessions first.
        try:
            session_id = proton.start(str(temp_file0.with_suffix("")), mode="pcsampling")
            proton.deactivate(session_id)
            temp_file1 = tmp_path / "test_profile1.hatchet"
            proton.start(str(temp_file1.with_suffix("")))
        except Exception as e:
            assert "Cannot add a session with the same profiler but a different mode than existing sessions" in str(e)
        finally:
            proton.finalize()


def test_profile_decorator(tmp_path: pathlib.Path):
    temp_file = tmp_path / "test_profile_decorator.hatchet"

    @proton.profile(name=str(temp_file.with_suffix("")))
    def foo0(a, b):
        return a + b

    foo0(1, 2)
    proton.finalize()
    assert temp_file.exists()

    @proton.profile
    def foo1(a, b):
        return a + b

    foo1(1, 2)
    proton.finalize()
    default_file = pathlib.Path(proton.DEFAULT_PROFILE_NAME + ".hatchet")
    assert default_file.exists()
    default_file.unlink()


def test_scope(tmp_path: pathlib.Path):
    # Scope can be annotated even when profiling is off
    with proton.scope("test"):
        pass

    temp_file = tmp_path / "test_scope.hatchet"
    proton.start(str(temp_file.with_suffix("")))
    with proton.scope("test"):
        pass

    @proton.scope("test")
    def foo():
        pass

    foo()

    proton.enter_scope("test")
    proton.exit_scope()

    proton.enter_scope("test0")
    proton.exit_scope("test0")

    proton.finalize()
    assert temp_file.exists()


def test_hook(tmp_path: pathlib.Path):
    temp_file = tmp_path / "test_hook.hatchet"
    session_id0 = proton.start(str(temp_file.with_suffix("")), hook="triton")
    proton.activate(session_id0)
    proton.activate(session_id0)
    assert len(
        HookManager.active_hooks) == 1, ("Activate a session multiple times should maintain a single instance of hook")
    assert list(HookManager.session_hooks[session_id0].values())[0] is True
    proton.deactivate(session_id0)
    assert list(HookManager.session_hooks[session_id0].values())[0] is False
    assert len(HookManager.active_hooks) == 0
    # Deactivate a session multiple times should not raise an error
    proton.deactivate(session_id0)
    proton.finalize(None)
    assert temp_file.exists()


def test_hook_manager(tmp_path: pathlib.Path):
    # Launch hook is a singleton
    HookManager.register(LaunchHook(), 0)
    HookManager.register(LaunchHook(), 0)
    assert len(HookManager.active_hooks) == 1
    assert isinstance(HookManager.active_hooks[0], LaunchHook)
    assert HookManager.session_hooks[0][HookManager.active_hooks[0]] is True

    # Only unregister one session
    HookManager.register(LaunchHook(), 1)
    HookManager.unregister(0)
    assert len(HookManager.active_hooks) == 1
    HookManager.unregister(1)
    assert len(HookManager.active_hooks) == 0

    # Heterogenous hooks
    HookManager.register(InstrumentationHook(""), 2)
    HookManager.register(LaunchHook(), 2)
    assert len(HookManager.active_hooks) == 2
    # Launch hook has a higher priority
    assert isinstance(HookManager.active_hooks[0], LaunchHook)
    assert isinstance(HookManager.active_hooks[1], InstrumentationHook)
    assert HookManager.session_hooks[2][HookManager.active_hooks[0]] is True
    assert HookManager.session_hooks[2][HookManager.active_hooks[1]] is True
    HookManager.unregister()
    assert len(HookManager.active_hooks) == 0


def test_scope_metrics(tmp_path: pathlib.Path):
    temp_file = tmp_path / "test_scope_metrics.hatchet"
    session_id = proton.start(str(temp_file.with_suffix("")))
    # Test different scope creation methods
    with proton.scope("test0", {"a": 1.0}):
        pass

    @proton.scope("test1", {"a": 1.0})
    def foo():
        pass

    foo()

    # After deactivation, the metrics should be ignored
    proton.deactivate(session_id)
    proton.enter_scope("test2", metrics={"a": 1.0})
    proton.exit_scope()

    # Metrics should be recorded again after reactivation
    proton.activate(session_id)
    proton.enter_scope("test3", metrics={"a": 1.0})
    proton.exit_scope()

    proton.enter_scope("test3", metrics={"a": 1.0})
    proton.exit_scope()

    # exit_scope can also take metrics
    proton.enter_scope("test4")
    proton.exit_scope(metrics={"b": 1.0})

    proton.finalize()

    assert temp_file.exists()
    with temp_file.open() as f:
        data = json.load(f)
    assert len(data[0]["children"]) == 4
    for child in data[0]["children"]:
        if child["frame"]["name"] == "test3":
            assert child["metrics"]["a"] == 2.0
        elif child["frame"]["name"] == "test4":
            assert child["metrics"]["b"] == 1.0


def test_scope_metrics_invalid(tmp_path: pathlib.Path):
    temp_file = tmp_path / "test_scope_metrics.hatchet"
    proton.start(str(temp_file.with_suffix("")))

    error = None

    try:
        with proton.scope("test0", {"a": 1.0}):
            pass

        with proton.scope("test0", {"a": 1}):
            pass
    except Exception as e:
        error = str(e)
    finally:
        proton.finalize()

    assert error is not None and "Metric value type mismatch for valueId 0 (a): current=double, new=int64_t" in error


def test_scope_properties(tmp_path: pathlib.Path):
    temp_file = tmp_path / "test_scope_properties.hatchet"
    proton.start(str(temp_file.with_suffix("")))
    # Properties do not aggregate
    proton.enter_scope("test0", metrics={"a (pty)": 1.0})
    proton.exit_scope()

    proton.enter_scope("test0", metrics={"a (pty)": 1.0})
    proton.exit_scope()

    proton.finalize()
    assert temp_file.exists()
    with temp_file.open() as f:
        data = json.load(f)
    for child in data[0]["children"]:
        if child["frame"]["name"] == "test0":
            assert child["metrics"]["a"] == 1.0


def test_scope_exclusive(tmp_path: pathlib.Path):
    temp_file = tmp_path / "test_scope_exclusive.hatchet"
    proton.start(str(temp_file.with_suffix("")))
    # metric a only appears in the outermost scope
    # metric b only appears in the innermost scope
    # both metrics do not appear in the root scope
    with proton.scope("test0", metrics={"a (exc)": 1}):
        with proton.scope("test1", metrics={"b (exc)": 1}):
            pass

    proton.finalize()
    assert temp_file.exists()
    with temp_file.open() as f:
        data = json.load(f)
    root_metrics = data[0]["metrics"]
    assert len(root_metrics) == 0
    test0_frame = data[0]["children"][0]
    test0_metrics = test0_frame["metrics"]
    assert len(test0_metrics) == 1
    assert test0_metrics["a"] == 1
    test1_frame = test0_frame["children"][0]
    test1_metrics = test1_frame["metrics"]
    assert len(test1_metrics) == 1
    assert test1_metrics["b"] == 1


def test_scope_vector_metrics(tmp_path: pathlib.Path):
    temp_file = tmp_path / "test_scope_vector_metrics.hatchet"
    proton.start(str(temp_file.with_suffix("")))
    proton.enter_scope("test0", metrics={"veci": [1, 2, 3], "vecd": [1.0, 2.0, 3.0]})
    proton.exit_scope()
    proton.finalize()
    assert temp_file.exists()
    with temp_file.open() as f:
        data = json.load(f)
    for child in data[0]["children"]:
        if child["frame"]["name"] == "test0":
            assert child["metrics"]["veci"] == [1, 2, 3]
            assert child["metrics"]["vecd"] == [1.0, 2.0, 3.0]


def test_state(tmp_path: pathlib.Path):
    temp_file = tmp_path / "test_state.hatchet"
    proton.start(str(temp_file.with_suffix("")))
    proton.enter_scope("test0")
    proton.enter_state("state")
    proton.enter_scope("test1", metrics={"a": 1.0})
    proton.exit_scope()
    proton.exit_state()
    proton.exit_scope()
    proton.finalize()
    assert temp_file.exists()
    with temp_file.open() as f:
        data = json.load(f)
    # test0->test1->state
    assert len(data[0]["children"]) == 1
    child = data[0]["children"][0]
    assert child["frame"]["name"] == "test0"
    assert len(child["children"]) == 1
    child = child["children"][0]
    assert child["frame"]["name"] == "test1"
    assert len(child["children"]) == 1
    child = child["children"][0]
    assert child["frame"]["name"] == "state"
    assert child["metrics"]["a"] == 1.0


def test_context_depth(tmp_path: pathlib.Path):
    temp_file = tmp_path / "test_context_depth.hatchet"
    session_id = proton.start(str(temp_file.with_suffix("")))
    assert proton.context.depth(session_id) == 0
    proton.enter_scope("test0")
    assert proton.context.depth(session_id) == 1
    proton.enter_scope("test1")
    assert proton.context.depth(session_id) == 2
    proton.exit_scope()
    assert proton.context.depth(session_id) == 1
    proton.exit_scope()
    assert proton.context.depth(session_id) == 0
    proton.finalize()


def test_throw(tmp_path: pathlib.Path):
    # Catch an exception thrown by c++
    session_id = 100
    temp_file = tmp_path / "test_throw.hatchet"
    activate_error = ""
    try:
        session_id = proton.start(str(temp_file.with_suffix("")))
        proton.activate(session_id + 1)
    except Exception as e:
        activate_error = str(e)
    finally:
        proton.finalize()
    assert "Session has not been initialized: " + str(session_id + 1) in activate_error

    deactivate_error = ""
    try:
        session_id = proton.start(str(temp_file.with_suffix("")))
        proton.deactivate(session_id + 1)
    except Exception as e:
        deactivate_error = str(e)
    finally:
        proton.finalize()
    assert "Session has not been initialized: " + str(session_id + 1) in deactivate_error


@pytest.mark.parametrize("disable", [True, False])
def test_profile_disable(disable, fresh_knobs, tmp_path: pathlib.Path):
    fresh_knobs.proton.disable = disable
    temp_file = tmp_path / "test_profile_disable.hatchet"
    proton.start(str(temp_file.with_suffix("")))
    proton.enter_scope("test0")
    proton.exit_scope()
    proton.finalize()
    if disable:
        assert not temp_file.exists()
    else:
        assert temp_file.exists()


def test_finalize_within_scope(tmp_path: pathlib.Path):
    temp_file = tmp_path / "test_finalize_within_scope.hatchet"
    session_id0 = proton.start(str(temp_file.with_suffix("")))
    with proton.scope("test0"):
        assert proton.context.depth(session_id0) == 1
        proton.finalize()
    assert temp_file.exists()
    temp_file1 = tmp_path / "test_finalize_within_scope1.hatchet"
    session_id1 = proton.start(str(temp_file1.with_suffix("")))
    depth = proton.context.depth(session_id1)
    assert depth == 0
    proton.finalize()


def test_data_api(tmp_path: pathlib.Path):
    temp_file = tmp_path / "test_data_api.hatchet"
    session_id = proton.start(str(temp_file.with_suffix("")))
    proton.enter_scope("test0")
    proton.exit_scope()
    proton.deactivate(session_id)
    json_data = proton.data.get(session_id)
    assert json_data is not None
    msgpack_data = proton.data.get_msgpack(session_id)
    assert isinstance(msgpack_data, bytes)
    is_complete = proton.data.is_phase_complete(session_id, 0)
    assert is_complete is False
    next_phase = proton.data.advance_phase(session_id)
    assert next_phase == 1
    is_complete = proton.data.is_phase_complete(session_id, 1)
    assert is_complete is False

    # Even if a phase has no GPU activity records, flushing should still mark it
    # as flushed.
    proton.activate(session_id)
    next_phase = proton.data.advance_phase(session_id)
    assert next_phase == 2
    proton.deactivate(session_id, flushing=True)
    assert proton.data.is_phase_complete(session_id, 1) is True
    assert proton.data.is_phase_complete(session_id, 2) is False

    # Test clear and clear_up_to_phase
    proton.data.clear(session_id, phase=0)
    proton.data.clear(session_id, phase=2, clear_up_to_phase=True)

    proton.finalize()
