# Copyright 2025-2026 XMOS LIMITED.
# This Software is subject to the terms of the XMOS Public Licence: Version 1.

from pathlib import Path
import subprocess
import re
import json

LINKING_EXE_RE = re.compile(r'.*Linking CXX executable\s*(\S+)')
TILE_HEADER_RE = re.compile(r'^Constraint check for tile\[(\d+)\]:')
MEM_LINE_RE = re.compile(r'^\s*Memory available:\s+(\d+),\s+used:\s+(\d+)\s+\.\s+(\w+)')
DETAIL_LINE_RE = re.compile(r'^\s*\(Stack:\s+(\d+),\s+Code:\s+(\d+),\s+Data:\s+(\d+)\)')


def parse_mem_report(stdout: str):
    current_tile = None
    cfg = None
    tile_of_interest = 1
    lines = stdout.splitlines()
    configs = {}
    for i, line in enumerate(lines):
        m_link = LINKING_EXE_RE.match(line)
        if m_link:
            exe = Path(m_link.group(1))
            cfg = exe.parent.name
            assert exe not in configs, f"{cfg} already present in configs. {configs}"
            configs[cfg] = {}
            current_tile = None
            continue
        m_tile = TILE_HEADER_RE.match(line)
        if m_tile:
            current_tile = int(m_tile.group(1))
            continue
        if current_tile == tile_of_interest:
            m_mem = MEM_LINE_RE.match(line)
            if m_mem and cfg:
                configs[cfg]["available"] = int(m_mem.group(1))
                configs[cfg]["used"] = int(m_mem.group(2))
                configs[cfg]["status"] = m_mem.group(3)
                continue
            m_det = DETAIL_LINE_RE.match(line)
            if m_det and cfg:
                configs[cfg]["stack"] = int(m_det.group(1))
                configs[cfg]["code"] = int(m_det.group(2))
                configs[cfg]["data"] = int(m_det.group(3))
                current_tile = None  # done with this tile
                cfg = None
    return configs

def write_rst_table(configs: dict, outfile: Path):
    lines = [
        ".. _mic_array_memory_usage:\n",
        ".. list-table:: Memory usage (in bytes)",
        "   :header-rows: 1",
        "   :widths: 8 8 8 6 6 6",
        "",
        "   * - Config",
        "     - Available on tile",
        "     - Used",
        "     - Stack",
        "     - Code",
        "     - Data",
    ]
    for cfg, data in configs.items():
        lines.append(f"   * - {cfg}")
        lines.append(f"     - {data.get('available','')}")
        lines.append(f"     - {data.get('used','')}")
        lines.append(f"     - {data.get('stack','')}")
        lines.append(f"     - {data.get('code','')}")
        lines.append(f"     - {data.get('data','')}")
    outfile.write_text("\n".join(lines))


def test_measure_memory(pytestconfig):
    """
    Compile and check the memory usage of barebones executables that:
        - include the mic using its default API
        - include it in the custom way by creating a mic_array::MicArray object

    Behaviour:
      • When --update is not passed:
          The memory numbers are checked against the reference JSON file
          (mic_array_memory.json). The test fails if any configuration deviates
          from the reference by more than 500 bytes. This ensures the mic array memory
          remain stable across changes.

      • When --update is passed:
          The reference JSON file is overwritten with the newly measured
          values. mic_array_memory_table.rst is also regenerated. No check
          against any reference values happens in this case

    Files:
      mic_array_memory.json         - stored reference memory values
      mic_array_memory_table.rst    - autogenerated RST table of results
    """
    update = pytestconfig.getoption("--update")
    pkg_dir = Path(__file__).parent
    subprocess_cwd = pkg_dir / "app_memory"
    build_dir = subprocess_cwd / "_build"

    subprocess.run(["cmake", "-G", "Unix Makefiles", "-B", build_dir], cwd=subprocess_cwd, capture_output=True, text=True, check=True)
    subprocess.run(["xmake", "-C", build_dir, "clean"], cwd=subprocess_cwd, capture_output=True, text=True, check=True)
    ret = subprocess.run(["xmake", "-C", build_dir], cwd=subprocess_cwd, capture_output=True, text=True, check=True)
    #print(f"STDOUT = {ret.stdout}")
    configs = parse_mem_report(ret.stdout)
    assert len(configs) > 0, f"ERROR: Couldn't parse memory report from stdout\n{ret.stdout}"

    json_out = pkg_dir / "mic_array_memory.json"
    if not update:
        with json_out.open("r") as f:
            ref_data = json.load(f)

    threshold = 500 # Allow up to 500 bytes change in overall memory before flagging
    for cfg, report in configs.items():
        print(f"cfg: {cfg}. report: {report}")
        assert report["status"] == "OKAY", f"ERROR: config {cfg} reports not OKAY status.\n{report}"
        if not update:
            assert cfg in ref_data, f"ERROR: config {cfg} not in reference yaml. Run test with pytest test_measure_memory.py --update to regenerate the reference yaml and rst"
            assert abs(report['used'] - ref_data[cfg]['used']) < threshold, (
                f"ERROR: Config {cfg}: Used memory {report['used']} off by more than "
                f"{threshold} bytes compared to the reference {ref_data[cfg]['used']}.\n"
                "If this is expected, run the test with pytest test_measure_memory.py --update "
                " to update the reference json and rst files.\n"
            )

    if update:
        with json_out.open("w") as f:
            json.dump(configs, f, indent=2)

        rst_out = pkg_dir / "mic_array_memory_table.rst"
        write_rst_table(configs, rst_out)
