"""Tests for memory/layout optimization passes (Issue #27)."""

from __future__ import annotations

from sagellm_backend.graph.graph import Graph, Node, OpType
from sagellm_backend.graph.passes.data_format_selection import DataFormatSelectionPass
from sagellm_backend.graph.passes.layout_conversion import LayoutConversionPass
from sagellm_backend.graph.passes.memory_reuse_analysis import MemoryReuseAnalysisPass


class TestLayoutConversionPass:
    """Tests for LayoutConversionPass."""

    def test_detect_nchw_to_nhwc_and_propagate(self) -> None:
        graph = Graph()
        x = Node(node_id="x", op_type=OpType.INPUT, attrs={"layout": "nchw"})
        t = Node(node_id="t", op_type=OpType.TRANSPOSE, attrs={"perm": [0, 2, 3, 1]})
        y = Node(node_id="y", op_type=OpType.MATMUL)

        for node in [x, t, y]:
            graph.add_node(node)
        graph.add_edge("x", "t")
        graph.add_edge("t", "y")
        graph.outputs = ["y"]

        result = LayoutConversionPass().run(graph)

        assert result.modified is True
        assert graph.nodes["t"].metadata["layout_conversion"] == "nchw_to_nhwc"
        assert graph.nodes["t"].metadata["inferred_layout"] == "nhwc"

    def test_detect_row_col_conversion(self) -> None:
        graph = Graph()
        a = Node(node_id="a", op_type=OpType.INPUT, attrs={"layout": "row_major"})
        t = Node(node_id="t", op_type=OpType.TRANSPOSE, attrs={"perm": [1, 0]})
        out = Node(node_id="out", op_type=OpType.OUTPUT)

        for node in [a, t, out]:
            graph.add_node(node)
        graph.add_edge("a", "t")
        graph.add_edge("t", "out")
        graph.outputs = ["out"]

        result = LayoutConversionPass().run(graph)

        assert result.modified is True
        assert graph.nodes["t"].metadata["layout_conversion"] == "row_major_to_col_major"


class TestDataFormatSelectionPass:
    """Tests for DataFormatSelectionPass."""

    def test_select_fp8_for_matmul_when_supported(self) -> None:
        graph = Graph(metadata={"supported_dtypes": ["float16", "bfloat16", "float8"]})
        mm = Node(node_id="mm", op_type=OpType.MATMUL, dtype="float32")
        graph.add_node(mm)
        graph.outputs = ["mm"]

        result = DataFormatSelectionPass().run(graph)

        assert result.modified is True
        assert graph.nodes["mm"].dtype == "float8"

    def test_select_bf16_for_norm_ops(self) -> None:
        graph = Graph(metadata={"supported_dtypes": ["float16", "bfloat16"]})
        ln = Node(node_id="ln", op_type=OpType.LAYER_NORM, dtype="float32")
        graph.add_node(ln)
        graph.outputs = ["ln"]

        result = DataFormatSelectionPass().run(graph)

        assert result.modified is True
        assert graph.nodes["ln"].dtype == "bfloat16"

    def test_force_dtype_unsupported_raises(self) -> None:
        graph = Graph(metadata={"supported_dtypes": ["float16"]})
        node = Node(
            node_id="x",
            op_type=OpType.ADD,
            dtype="float32",
            attrs={"force_dtype": "float8"},
        )
        graph.add_node(node)
        graph.outputs = ["x"]

        try:
            DataFormatSelectionPass().run(graph)
        except ValueError as exc:
            assert "unsupported dtype" in str(exc)
        else:
            raise AssertionError("Expected ValueError for unsupported force_dtype")


class TestMemoryReuseAnalysisPass:
    """Tests for MemoryReuseAnalysisPass."""

    def test_lifecycle_and_pool_slot_planning(self) -> None:
        graph = Graph()
        x = Node(node_id="x", op_type=OpType.INPUT)
        a = Node(node_id="a", op_type=OpType.ADD, attrs={"size_bytes": 4096})
        b = Node(node_id="b", op_type=OpType.MUL, attrs={"size_bytes": 4096})
        out = Node(node_id="out", op_type=OpType.OUTPUT)

        for node in [x, a, b, out]:
            graph.add_node(node)
        graph.add_edge("x", "a")
        graph.add_edge("a", "b")
        graph.add_edge("b", "out")
        graph.outputs = ["out"]

        result = MemoryReuseAnalysisPass().run(graph)

        assert result.modified is True
        assert "memory_plan" in graph.metadata
        assert graph.metadata["memory_plan"]["planned_allocations"] >= 1
        assert "memory_pool_slot" in graph.nodes["a"].metadata
        assert "memory_pool_slot" in graph.nodes["b"].metadata

    def test_cpu_validation_end_to_end(self) -> None:
        graph = Graph(metadata={"supported_dtypes": ["float16", "float8"]})
        x = Node(node_id="x", op_type=OpType.INPUT, device="cpu", attrs={"layout": "nchw"})
        t = Node(
            node_id="t",
            op_type=OpType.TRANSPOSE,
            device="cpu",
            attrs={"perm": [0, 2, 3, 1], "shape": [1, 64, 64, 32]},
            dtype="float32",
        )
        mm = Node(
            node_id="mm",
            op_type=OpType.MATMUL,
            device="cpu",
            attrs={"shape": [1, 64, 128], "size_bytes": 16384},
            dtype="float32",
        )
        out = Node(node_id="out", op_type=OpType.OUTPUT, device="cpu")

        for node in [x, t, mm, out]:
            graph.add_node(node)

        graph.add_edge("x", "t")
        graph.add_edge("t", "mm")
        graph.add_edge("mm", "out")
        graph.outputs = ["out"]

        layout_result = LayoutConversionPass().run(graph)
        format_result = DataFormatSelectionPass().run(graph)
        memory_result = MemoryReuseAnalysisPass().run(graph)

        assert layout_result.modified is True
        assert format_result.modified is True
        assert memory_result.modified is True
        assert graph.nodes["t"].metadata["layout_conversion"] == "nchw_to_nhwc"
        assert graph.nodes["mm"].dtype == "float8"
        assert graph.metadata["memory_plan"]["planned_allocations"] >= 1
