from __future__ import annotations

import hashlib
import xml.etree.ElementTree as ET
from pathlib import Path

import pytest

from scripts import build_mlir_ir_graphs as builder
from scripts import mlir_graph_common as common


def digest(path: Path) -> str:
    value = hashlib.sha256()
    value.update(path.read_bytes())
    return value.hexdigest()


def write_mlir(tmp_path: Path, text: str, name: str = "fixture.mlir") -> Path:
    path = tmp_path / name
    path.write_text(text.strip() + "\n", encoding="utf-8")
    return path


def test_onnx_ssa_dependency_and_static_order(tmp_path: Path) -> None:
    path = write_mlir(
        tmp_path,
        """
module {
  func.func @main(%arg0: tensor<1xf32>) -> tensor<1xf32> {
    %relu = "onnx.Relu"(%arg0) {onnx_node_name = "Relu_0"} : (tensor<1xf32>) -> tensor<1xf32>
    return %relu : tensor<1xf32>
  }
}
""",
    )
    parsed = common.parse_mlir(path, "TEST:fp32:ONNX")
    relu = next(row for row in parsed.operations if row["operation"] == "onnx.Relu")
    returned = next(row for row in parsed.operations if row["operation"] == "func.return")
    assert relu["results"] == ["%relu"]
    assert returned["operands"] == ["%relu"]
    assert parsed.unresolved_use_count == 0
    assert parsed.duplicate_definition_count == 0
    assert parsed.producer_after_consumer_count == 0

    svg, metadata = common.render_execution_dependency_svg(
        parsed,
        title="ONNX primary fixture",
        graph_id="TEST:fp32:ONNX",
        stage="ONNX",
    )
    ET.fromstring(svg)
    assert metadata["rendered_operation_nodes"] == len(parsed.operations)
    assert "STATIC_MLIR_PROGRAM_ORDER" in svg


def test_grouped_results_expand_and_resolve(tmp_path: Path) -> None:
    path = write_mlir(
        tmp_path,
        """
module {
  func.func @main(%arg0: tensor<4xf32>) {
    %packed : 2 = "onnx.Split"(%arg0) : (tensor<4xf32>) -> (tensor<2xf32>, tensor<2xf32>)
    "onnx.Add"(%packed#0, %packed#1) : (tensor<2xf32>, tensor<2xf32>) -> ()
    return
  }
}
""",
    )
    parsed = common.parse_mlir(path, "GROUPED:public_quantized:ONNX")
    split = next(row for row in parsed.operations if row["operation"] == "onnx.Split")
    assert split["results"] == ["%packed#0", "%packed#1"]
    assert parsed.unresolved_use_count == 0


def test_primary_only_selection_is_exactly_21_pairs_42_graphs(tmp_path: Path) -> None:
    model_ids = [f"M{index:02d}" for index in range(21)]
    rows: list[dict[str, str]] = []
    for model_id in model_ids:
        for variant in common.VARIANTS:
            directory = tmp_path / "models" / model_id / variant
            directory.mkdir(parents=True)
            onnx_path = directory / "onnx.mlir"
            onnx_path.write_text("module {}\n", encoding="utf-8")
            rows.append(
                {
                    "model_id": model_id,
                    "task": "fixture",
                    "variant": variant,
                    "onnx_status": "PASS",
                    "onnx_artifact": str(onnx_path.relative_to(tmp_path)),
                    "onnx_sha256": digest(onnx_path),
                    "affine_scf_memref_status": "FAIL",
                    "affine_scf_memref_artifact": "NONE",
                    "affine_scf_memref_sha256": "NONE",
                    "krnl_status": "FAIL",
                    "krnl_artifact": "NONE",
                    "krnl_sha256": "NONE",
                    "llvm_status": "FAIL",
                    "llvm_artifact": "NONE",
                    "llvm_sha256": "NONE",
                    "last_fully_successful_ir": "ONNX",
                }
            )

    by_key = builder.require_matrix_rows(rows)
    specs, pairs = builder.build_specs(by_key, tmp_path, primary_only=True)
    assert len(specs) == 42
    assert len({row["graph_id"] for row in specs}) == 42
    assert {row["stage"] for row in specs} == {"ONNX"}
    assert {row["graph_role"] for row in specs} == {"PRIMARY_ONNX_DIALECT"}
    assert len(pairs) == 21
    assert {row["pair_common_stage"] for row in pairs} == {"ONNX"}
    assert {row["supplemental_lower_graphs"] for row in pairs} == {
        "0 (primary ONNX reused as common view)"
    }
    with pytest.raises(ValueError, match="duplicate IR coverage row"):
        builder.require_matrix_rows([*rows, dict(rows[0])])