"""A memory tester for a multiprocessed program.

Returns a namedtuple of first, max, last memory.

It takes readings of the main and child processes at intervals and reports the entry,
exit, and max memory usage during the running of a function.

Used as a decorator for a function via @test_memory
"""

from __future__ import annotations

import multiprocessing
import os
import threading
import time
from typing import TYPE_CHECKING, NamedTuple

import psutil

from .returnable_thread import ReturnableThread

if TYPE_CHECKING:
    from collections.abc import Callable


class ResourceResult(NamedTuple):
    """Hold a memory result with the members as strings like '023.456GB'."""

    first: float
    max: float
    last: float

    def __str__(self):
        # Most numbers will be between 0 and N cores, which is less than 1000 for cpus
        # (usually) and 0 and 1024 for memory, but allow for larger numbers
        f, m, l = [  # noqa: E741
            f"{s:.3g}" if s < 1000 else f"{s:.0f}"  # noqa: PLR2004
            for s in (self.first, self.max, self.last)
        ]
        return f"[{f}=>({m})=>{l}]"


def get_process(pid: int | None) -> psutil.Process | None:
    """Return the process based on PID or None."""
    try:
        return psutil.Process(pid)
    except psutil.NoSuchProcess:
        return None


def get_cpu_cores(
    process: psutil.Process | None = None,
) -> tuple[float, psutil.Process]:
    """Get this process' cpu percentage along with child usage.

    Note, first call will return 0
    """
    if process is None:
        process = psutil.Process(os.getpid())
    cpu_pct = process.cpu_percent()
    for child in multiprocessing.active_children():
        child_process = get_process(child.pid)
        if child_process is not None:
            cpu_pct += child_process.cpu_percent()
    return cpu_pct / 100, process


def get_memory_mb(
    process: psutil.Process | None = None,
) -> tuple[float, psutil.Process]:
    """Get this process' megabytes of reported memory along with child usage."""
    if process is None:
        process = psutil.Process(os.getpid())
    gb_mem = process.memory_info().rss / 1024**2
    for child in multiprocessing.active_children():
        child_process = get_process(child.pid)
        if child_process is not None:
            gb_mem += child_process.memory_info().rss / 1024**2
    return gb_mem, process


def resource_worker(
    exit_event: threading.Event,
    delay_s: float = 0.5,
) -> tuple[ResourceResult, ResourceResult]:
    """Check all memory usage while event.is_set()."""
    cpu_first, process = get_cpu_cores()
    mem_first, process = get_memory_mb(process)
    mem_max: float = mem_first
    cpu_max: float = cpu_first
    while True:
        if exit_event.wait(delay_s):
            cpu = ResourceResult(cpu_first, cpu_max, get_cpu_cores(process)[0])
            mem = ResourceResult(mem_first, mem_max, get_memory_mb(process)[0])
            return mem, cpu
        mem_max = max(get_memory_mb(process)[0], mem_max)
        cpu_max = max(get_cpu_cores(process)[0], cpu_max)


# def mb_to_str(mb: float) -> str:
#     """Convert MB to KB or GB as needed and tag"""
#     if mb < 1:
#         return f"{mb*1024:07.3f}KB"
#     if mb >= 1024:
#         return f"{mb/1024:07.3f}GB"
#     return f"{mb:07.3}MB"


def hhmmss(s: float) -> str:
    """Convert decimal seconds to hh:mm:ss.sss.

    Internal function for simple printing
    """
    ms = s % 1
    s = int(s - ms)
    ms = round(ms * 1000)
    ss = s % 60
    s = int((s - ss) / 60)
    mm = s % 60
    hh = int((s - mm) / 60)
    pp = f"{ss:02d}.{ms:03d}"
    key = "ss.sss"
    if hh > 0 or mm > 0:
        pp = f"{mm:02d}:{pp}"
        key = f"mm:{key}"
    if hh > 0:
        pp = f"{hh}:{pp}"
        key = f"hh:{key}"
    return f"{pp} ({key})"


def test_memory(func: Callable) -> Callable:
    """Decorator to print memory usage during a function.

    Includes all child processes of the program.

    Decorate the function with @test_memory.

    Memory will be printed in gigabytes for entry, max, and exit values.
    """

    def wrapper(*args, **kwargs):  # noqa: ANN202
        exit_event = threading.Event()
        funcname = func.__name__
        rt = ReturnableThread[tuple[ResourceResult, ResourceResult]](
            target=resource_worker,
            args=(exit_event,),
        )
        rt.start()
        t1 = time.time()

        ret = func(*args, **kwargs)

        t2 = time.time()
        exit_event.clear()
        exit_event.set()
        mem, cpu = rt.join()
        units = ["MB", "GB", "TB", "PB"]  # future proof for a while
        unit = 0
        while any(m > 1024 for m in mem):  # noqa: PLR2004
            mem = ResourceResult(mem.first / 1024, mem.max / 1024, mem.last / 1024)
            unit += 1
        m_unit = units[unit]
        print(f"*** ({funcname}) {hhmmss(t2 - t1)} *** {mem}{m_unit} {cpu}cores")  # noqa: T201
        return ret

    return wrapper
