# SPDX-FileCopyrightText: 2025 MetaMachines LLC
#
# SPDX-License-Identifier: MIT

from __future__ import annotations

import ctypes
import os
import sys
import unittest
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[1]
SRC_DIR = REPO_ROOT / "src"
EXAMPLES_DIR = REPO_ROOT / "examples"
for path in (SRC_DIR, EXAMPLES_DIR):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))


from mm_ptx import get_include_dir, get_ptx_inject_header
from mm_ptx.ptx_inject import PTXInject, MutType
import mm_ptx.stack_ptx as stack_ptx
from stack_ptx_default_types import Stack, PtxInstruction
from stack_ptx_default_types import compiler as stack_ptx_compiler


STACK_PTX_KWARGS = {
    "execution_limit": 100,
    "max_ast_size": 100,
    "max_ast_to_visit_stack_depth": 20,
    "stack_size": 128,
    "max_frame_depth": 4,
    "store_size": 16,
}


PTX_SAMPLE = (
    "// pre\n"
    "// PTX_INJECT_START func\n"
    "// _x0 i f32 F32 x\n"
    "// _x1 m f32 F32 y\n"
    "// _x2 o f32 F32 z\n"
    "// PTX_INJECT_END\n"
    "// post\n"
)

PTX_SAMPLE_TWO = (
    "// pre\n"
    "// PTX_INJECT_START func_a\n"
    "// _x0 i f32 F32 x\n"
    "// _x1 o f32 F32 y\n"
    "// PTX_INJECT_END\n"
    "// mid\n"
    "// PTX_INJECT_START func_b\n"
    "// _x0 i f32 F32 x\n"
    "// _x1 o f32 F32 y\n"
    "// PTX_INJECT_END\n"
    "// post\n"
)


def _compile_stub_xy_to_yz(func) -> str:
    registry = stack_ptx.RegisterRegistry()
    registry.add(func["x"].reg, Stack.f32, name="x")
    registry.add(func["y"].reg, Stack.f32, name="y")
    registry.add(func["z"].reg, Stack.f32, name="z")
    registry.freeze()

    instructions = [
        registry.x,
        registry.y,
        PtxInstruction.add_ftz_f32,
        Stack.f32.dup,
        registry.x,
        PtxInstruction.add_ftz_f32,
    ]

    stub = stack_ptx_compiler.compile(
        registry=registry,
        instructions=instructions,
        requests=[registry.z, registry.y],
        **STACK_PTX_KWARGS,
    )
    return stub.rstrip("\x00")


def _compile_stub_xx_to_y(func, op) -> str:
    registry = stack_ptx.RegisterRegistry()
    registry.add(func["x"].reg, Stack.f32, name="x")
    registry.add(func["y"].reg, Stack.f32, name="y")
    registry.freeze()

    instructions = [
        registry.x,
        registry.x,
        op,
    ]

    stub = stack_ptx_compiler.compile(
        registry=registry,
        instructions=instructions,
        requests=[registry.y],
        **STACK_PTX_KWARGS,
    )
    return stub.rstrip("\x00")


class TestStackPtxInjectParsing(unittest.TestCase):
    def test_parse_and_render(self):
        with PTXInject(PTX_SAMPLE) as inject:
            self.assertEqual(len(inject.injects), 1)
            func = inject["func"]

            self.assertEqual(func["x"].mut_type, MutType.IN)
            self.assertEqual(func["x"].data_type, "F32")
            self.assertEqual(func["x"].register_type, "f32")
            self.assertEqual(func["x"].reg, "_x0")

            self.assertEqual(func["y"].mut_type, MutType.MOD)
            self.assertEqual(func["y"].data_type, "F32")
            self.assertEqual(func["y"].register_type, "f32")
            self.assertEqual(func["y"].reg, "_x1")

            self.assertEqual(func["z"].mut_type, MutType.OUT)
            self.assertEqual(func["z"].data_type, "F32")
            self.assertEqual(func["z"].register_type, "f32")
            self.assertEqual(func["z"].reg, "_x2")

            ptx_stub = _compile_stub_xy_to_yz(func)
            rendered = inject.render_ptx({"func": ptx_stub})

        self.assertIn("// pre", rendered)
        self.assertIn("// post", rendered)
        self.assertIn("add.ftz.f32", rendered)
        self.assertIn("mov.f32 %_x2", rendered)
        self.assertIn("mov.f32 %_x1", rendered)
        self.assertNotIn("PTX_INJECT_START", rendered)
        self.assertNotIn("PTX_INJECT_END", rendered)


class TestStackPtxInjectMultiSite(unittest.TestCase):
    def test_render_two_sites(self):
        with PTXInject(PTX_SAMPLE_TWO) as inject:
            self.assertEqual(len(inject.injects), 2)
            self.assertEqual(inject.injects[0].name, "func_a")
            self.assertEqual(inject.injects[1].name, "func_b")

            stub_a = _compile_stub_xx_to_y(inject["func_a"], PtxInstruction.add_ftz_f32)
            stub_b = _compile_stub_xx_to_y(inject["func_b"], PtxInstruction.mul_ftz_f32)
            rendered = inject.render_ptx({"func_a": stub_a, "func_b": stub_b})

        self.assertIn("// pre", rendered)
        self.assertIn("// mid", rendered)
        self.assertIn("// post", rendered)
        self.assertIn(stub_a, rendered)
        self.assertIn(stub_b, rendered)
        self.assertLess(rendered.index(stub_a), rendered.index(stub_b))


class TestStackPtxInjectCudaIntegration(unittest.TestCase):
    def test_cuda_output_value(self):
        if os.getenv("MM_PTX_RUN_CUDA_TESTS") != "1":
            self.skipTest("Set MM_PTX_RUN_CUDA_TESTS=1 to run CUDA integration tests.")

        try:
            from cuda.core import Device, Program, ProgramOptions, LaunchConfig, launch, DeviceMemoryResource
        except Exception as exc:
            self.skipTest(f"cuda.core not available: {exc}")

        try:
            from cuda.bindings import runtime as cudart
        except Exception as exc:
            self.skipTest(f"cuda.bindings.runtime not available: {exc}")

        try:
            dev = Device()
            dev.set_current()
        except Exception as exc:
            self.skipTest(f"CUDA device not available: {exc}")

        def _make_program_options(**kwargs):
            include_dirs = kwargs.pop("include_dirs", None)
            if include_dirs:
                for include_key in ("include_path", "include_paths", "include_dirs"):
                    try:
                        return ProgramOptions(**kwargs, **{include_key: include_dirs})
                    except TypeError:
                        continue
            return ProgramOptions(**kwargs)

        capability = dev.compute_capability
        arch = f"sm_{capability.major}{capability.minor}"

        header_path = get_ptx_inject_header().replace("\\", "/")
        cuda_code = f"""
#include "{header_path}"

extern "C"
__global__
void
kernel(float* out) {{
    float x = 5.0f;
    float y = 3.0f;
    float z = 0.0f;
    for (int i = 0; i < 2; i++) {{
        PTX_INJECT("func",
            PTX_IN (F32, x, x),
            PTX_MOD(F32, y, y),
            PTX_OUT(F32, z, z)
        );
    }}
    out[0] = z;
}}
"""

        program_options = _make_program_options(
            std="c++11",
            arch=arch,
            include_dirs=[get_include_dir()],
        )
        prog = Program(cuda_code, code_type="c++", options=program_options)
        mod = prog.compile("ptx", logs=sys.stdout)
        annotated_ptx = mod.code.decode("utf-8")

        with PTXInject(annotated_ptx) as inject:
            ptx_stub = _compile_stub_xy_to_yz(inject["func"])
            rendered_ptx = inject.render_ptx({"func": ptx_stub})

        prog = Program(rendered_ptx, code_type="ptx", options=ProgramOptions(arch=arch))
        mod = prog.compile("cubin", logs=sys.stdout)
        ker = mod.get_kernel("kernel")

        dmem = DeviceMemoryResource(dev)
        out_buf = dmem.allocate(ctypes.sizeof(ctypes.c_float))

        stream = dev.default_stream
        launch(stream, LaunchConfig(grid=1, block=1), ker, out_buf)
        stream.sync()

        out_host = ctypes.c_float(0.0)
        out_host_ptr = ctypes.addressof(out_host)
        out_ptr = out_buf.handle
        if hasattr(out_ptr, "getPtr"):
            out_ptr = out_ptr.getPtr()
        out_ptr = int(out_ptr)

        err, = cudart.cudaMemcpy(
            out_host_ptr,
            out_ptr,
            ctypes.sizeof(ctypes.c_float),
            cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
        )
        if err != cudart.cudaError_t.cudaSuccess:
            self.fail(f"cudaMemcpy failed: {err}")

        self.assertAlmostEqual(out_host.value, 18.0, places=3)
