import argparse
import copy
import json
from pathlib import Path
import subprocess
import tempfile


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--compiler", required=True)
    args = parser.parse_args()
    repo = Path(__file__).resolve().parents[2]
    config_path = repo / "validation/pimsim_configs/pimcomp/arch-a/latency_config.json"
    config = json.loads(config_path.read_text())
    network = config["chip_config"]["network_config"]
    network["net_config_file_path"] = str(config_path.parent / network["net_config_file_path"])
    with tempfile.TemporaryDirectory(prefix="raptor-target-config-") as temporary:
        root = Path(temporary)
        for name, memory, valid in (
            ("absent", None, True),
            ("defaults", {}, True),
            ("explicit", config["chip_config"]["global_memory_config"], True),
            ("zero-service", {"read_latency_cycle": 0, "write_latency_cycle": 0}, True),
            ("bad-object", [], False),
            ("bad-width", {"data_width": 0}, False),
            ("bad-period", {"period": -1}, False),
        ):
            candidate = copy.deepcopy(config)
            if memory is None:
                del candidate["chip_config"]["global_memory_config"]
            else:
                candidate["chip_config"]["global_memory_config"] = memory
            path = root / f"{name}.json"
            path.write_text(json.dumps(candidate))
            command = [args.compiler,
                str(repo / "validation/operations/relu/after_conv/relu_after_conv.onnx"),
                "--maccel=PIM", "--EmitSpatial", "--crossbar-size=128",
                "--crossbar-count=96", "--core-count=168", "--pipeline=1",
                "--mlir-disable-threading", f"--pim-target-config={path}",
                "-o", str(root / name)]
            result = subprocess.run(command, stdout=subprocess.PIPE,
                                    stderr=subprocess.STDOUT, text=True)
            assert (result.returncode == 0) == valid, (name, result.stdout)
            if not valid:
                assert "Pim target config" in result.stdout, result.stdout


if __name__ == "__main__":
    main()