import re
import shlex
import subprocess
from pathlib import Path
from colorama import Fore, Style
from subprocess_utils import run_command_with_reporter

PIM_PASS_LABELS = (
    ("ONNXToSpatialPass", "ONNX to Spatial"),
    ("MergeComputeNodesPass", "Merge Compute Nodes"),
    ("SpatialToPimPass", "Spatial to PIM"),
    ("PimBufferizationPass", "Bufferize PIM"),
    ("HostConstantFoldingPass", "Fold Host Constants"),
    ("PimMemoryCoalescingPass", "Coalesce Local Memory"),
    ("PimLocalMemoryPlanningPass", "Plan Local Memory"),
    ("VerificationPass", "Verify PIM"),
    ("EmitPimCodePass", "Emit PIM Code"),
)
PIM_PASS_LABEL_BY_SUFFIX = dict(PIM_PASS_LABELS)
TIMING_LINE_RE = re.compile(r"^\s*([0-9]+\.[0-9]+)\s+\(\s*[0-9.]+%\)\s+(.+?)\s*$")


def _parse_pim_pass_timings(output_text):
    pass_timings = {}
    for line in output_text.splitlines():
        match = TIMING_LINE_RE.match(line)
        if not match:
            continue

        duration = float(match.group(1))
        pass_name = match.group(2)
        for suffix, label in PIM_PASS_LABEL_BY_SUFFIX.items():
            if pass_name.endswith(suffix):
                pass_timings[label] = pass_timings.get(label, 0.0) + duration
                break

    return pass_timings


def _format_command(cmd):
    return shlex.join(str(arg) for arg in cmd)


def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
                        crossbar_size, crossbar_count, core_count=None,
                        pim_memory_report="none", raptor_extra_args=None, cwd=None, verbose=False,
                        reporter=None, timeout_sec=None):
    # Define the arguments, with the possibility to set crossbar size and count
    args = [
        network_path,
        "-o",
        output_base,
        "--maccel=PIM",
        "--EmitPimCodegen",
        f"--crossbar-size={crossbar_size}",
        f"--crossbar-count={crossbar_count}",
    ]
    if core_count is not None:
        args.append(f"--core-count={core_count}")
    if pim_memory_report != "none":
        args.append(f"--pim-memory-report={pim_memory_report}")
    if raptor_extra_args:
        args.extend(str(arg) for arg in raptor_extra_args)
    if verbose:
        args.append("--enable-timing")

    cmd = [str(raptor_onnx_path)] + [str(arg) for arg in args]
    if reporter is not None:
        reporter.log(f"  Raptor command: {_format_command(cmd)}")
    else:
        print(f"Raptor command: {_format_command(cmd)}")

    try:
        output_text = run_command_with_reporter(
            cmd,
            cwd=cwd,
            reporter=reporter,
            capture_output=True,
            timeout_sec=timeout_sec,
        )
        if reporter is None:
            print(Fore.GREEN + "Raptor execution successful" + Style.RESET_ALL)
        return _parse_pim_pass_timings(output_text)
    except subprocess.CalledProcessError:
        if reporter is None:
            print(Fore.RED + "Raptor execution failed" + Style.RESET_ALL)
        raise