"""Intermediate emission must stop before unsupported backend transformations."""
import argparse
import json
import os
from pathlib import Path
import subprocess
import tempfile

import onnx
from onnx import TensorProto, helper

REPO = Path(__file__).resolve().parents[2]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--compiler", required=True)
    args = parser.parse_args()
    with tempfile.TemporaryDirectory(prefix="raptor-emission-") as temporary:
        root = Path(temporary)

        def model(name, elements):
            tensors = [helper.make_tensor_value_info(n, TensorProto.FLOAT, [1, elements])
                       for n in ("x", "y", "z")]
            graph = helper.make_graph([helper.make_node("Add", ["x", "y"], ["z"])],
                                      name, tensors[:2], tensors[2:])
            value = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
            value.ir_version = 10
            path = root / f"{name}.onnx"
            onnx.save(value, path)
            return path

        small = model("small", 16)
        # This static type is legal in Spatial. PIM's signed byte-size operand cannot encode it.
        large = model("large", 536870912)
        for source, targets in ((small, ("EmitSpatial", "EmitPim", "EmitPimBufferized", "EmitPimCodegen")),
                                (large, ("EmitSpatial", "EmitPim"))):
            for target in targets:
                out = root / source.stem / target
                out.mkdir(parents=True)
                result = subprocess.run([args.compiler, str(source), "-o", str(out / "model"),
                    "--maccel=PIM", f"--{target}", "--pipeline=2", "--core-count=168",
                    "--crossbar-count=96", "--crossbar-size=128", "--mlir-disable-threading", "--pim-dump-ir",
                    f"--pim-target-config={REPO}/validation/pimsim_configs/pimcomp/arch-a/throughput_config_1000ms.json"],
                    capture_output=True, text=True, timeout=30, env=dict(os.environ, OMP_NUM_THREADS="1"))
                output = result.stdout + result.stderr
                if source == large and target == "EmitPim":
                    assert result.returncode != 0 and "host input copy byte size is outside representable range" in output, output
                    continue
                assert result.returncode == 0, output
                if target == "EmitPimCodegen":
                    assert (out / "pim/config.json").exists()
                    report = json.loads((out / "reports/pipeline_selection.json").read_text())
                    assert report["generated_candidates"] == 1
                    assert report["single_candidate_analysis_skipped"]
                    assert report["lowered_candidates"] == 1
                    assert "no cost comparison required" in report["reason"]
                else:
                    assert "intermediate emission retains the original schedule" in output
                    assert not (out / "reports/pipeline_selection.json").exists()
                    assert not (out / "reports/pipeline_periodic.json").exists()
                    assert not (out / "pim").exists()
                    name = {"EmitSpatial": "spatial4_scheduled", "EmitPim": "pim0",
                            "EmitPimBufferized": "pim1_buff"}[target]
                    ir = (out / f"dialects/{name}.mlir").read_text()
                    if target == "EmitSpatial":
                        assert "spat.scheduled_compute" in ir and "pim.core" not in ir
                        assert not (out / "dialects/pim0.mlir").exists()
                    elif target == "EmitPim":
                        assert "pim.core" in ir and "tensor<" in ir
                        assert not (out / "dialects/pim1_buff.mlir").exists()
                    else:
                        assert "pim.core" in ir and "memref<" in ir
        out = root / "explicit-diagnostic"
        out.mkdir()
        result = subprocess.run([args.compiler, str(small), "-o", str(out / "model"),
            "--maccel=PIM", "--EmitSpatial", "--pipeline=2", "--core-count=168",
            "--crossbar-count=96", "--crossbar-size=128", "--mlir-disable-threading",
            "--pim-periodic-diagnostics",
            f"--pim-target-config={REPO}/validation/pimsim_configs/pimcomp/arch-a/throughput_config_1000ms.json"],
            capture_output=True, text=True, timeout=30)
        assert result.returncode == 0, result.stdout + result.stderr
        assert (out / "reports/pipeline_periodic.json").exists()
        report = json.loads((out / "reports/pipeline_periodic.json").read_text())
        assert report["diagnostic_only"]
        assert not report["selection_eligible"]
        assert not (out / "pim").exists()


if __name__ == "__main__":
    main()