import os
import psutil
import argparse


def check_memery_maps(process):
    memory_maps = process.memory_maps()
    memory_maps.sort(key=lambda x: x.rss, reverse=True)

    print(f"PID: {process.pid} Memory maps:")
    # 添加表头
    print(f"{'Idx':<3}  {'RSS (MB)':<12} {'Path'}")
    print("-" * 50)

    for idx, item in enumerate(memory_maps):
        path = item.path
        rss_mb = item.rss / 1024 / 1024  # byte to mb
        print(f"{idx:<3}  {rss_mb:<12.2f} {path}")

    total_rss = sum(item.rss for item in memory_maps)
    print(f"Total RSS: {total_rss / 1024 / 1024:.2f} MB")


def check_open_file(process):
    open_files = process.open_files()
    print(f"PID: {process.pid} Open files[{len(open_files)}]:")
    print(f"{'FD':<8} {'Path'}")
    print("-" * 50)
    for item in open_files:
        print(f"{item.fd:<8} {item.path}")


def check_thread(process):
    print(f"PID: {process.pid} Threads[{len(process.threads())}]:")
    for idx, th in enumerate(process.threads()):
        print(idx + 1, th)


def main(pid: int = None):
    process = psutil.Process(pid or os.getpid())
    check_memery_maps(process)
    print("=" * 50)
    check_open_file(process)
    print("=" * 50)
    check_thread(process)


if __name__ == "__main__":
    """
    python -m huibiao_framework.utils.process_check
    """
    parser = argparse.ArgumentParser(
        description="Check system memory maps and open files for a process"
    )
    parser.add_argument(
        "pid",
        nargs="?",
        type=int,
        help="Process ID to check (default: current process)",
    )

    args = parser.parse_args()
    main(args.pid)
