"""Parse python -X importtime output and produce a sorted breakdown."""

import os
import re
import subprocess
import sys


def parse_importtime(stderr_lines):
    pattern = re.compile(
        r"import time:\s+(\d+)\s+\|\s+(\d+)\s+\|\s+(\s*)([\w.]+)"
    )
    results = []
    for line in stderr_lines:
        m = pattern.match(line)
        if m:
            self_us = int(m.group(1))
            cumul_us = int(m.group(2))
            indent = len(m.group(3)) // 2
            module = m.group(4)
            results.append((module, self_us, cumul_us, indent))
    return results


def main():
    target = sys.argv[1] if len(sys.argv) > 1 else "import rich"
    venv_python = os.path.expanduser("~/rich/.venv/bin/python")

    proc = subprocess.run(
        [venv_python, "-X", "importtime", "-c", target],
        capture_output=True,
        text=True,
    )
    entries = parse_importtime(proc.stderr.splitlines())
    entries.sort(key=lambda e: e[1], reverse=True)

    print(f"{'Module':<50} {'Self (us)':>12} {'Cumul (us)':>12} {'Depth':>6}")
    print("-" * 82)
    for mod, self_us, cumul_us, depth in entries[:40]:
        print(f"{mod:<50} {self_us:>12,} {cumul_us:>12,} {depth:>6}")

    if len(sys.argv) > 2:
        with open(sys.argv[2], "w") as f:
            f.write("module\tself_us\tcumul_us\tdepth\n")
            for mod, self_us, cumul_us, depth in entries:
                f.write(f"{mod}\t{self_us}\t{cumul_us}\t{depth}\n")
        print(f"\nTSV written to {sys.argv[2]}")


if __name__ == "__main__":
    main()
