from __future__ import annotations

import sys

from cupy.cuda import memory_hook


class DebugPrintHook(memory_hook.MemoryHook):
    """Memory hook that prints debug information.

    This memory hook outputs the debug information of input arguments of
    ``malloc`` and ``free`` methods involved in the hooked functions
    at postprocessing time (that is, just after each method is called).

    Example:
        The basic usage is to use it with ``with`` statement.

        Code example::

            >>> import cupy
            >>> from cupy.cuda import memory_hooks
            >>>
            >>> cupy.cuda.set_allocator(cupy.cuda.MemoryPool().malloc)
            >>> with memory_hooks.DebugPrintHook():
            ...     x = cupy.array([1, 2, 3])
            ...     del x  # doctest:+SKIP

        Output example::

            {"hook":"alloc","device_id":0,"mem_size":512,"mem_ptr":150496608256}
            {"hook":"malloc","device_id":0,"size":24,"mem_size":512,"mem_ptr":150496608256,"pmem_id":"0x7f39200c5278"}
            {"hook":"free","device_id":0,"mem_size":512,"mem_ptr":150496608256,"pmem_id":"0x7f39200c5278"}

        where the output format is JSONL (JSON Lines) and
        ``hook`` is the name of hook point, and
        ``device_id`` is the CUDA Device ID, and
        ``size`` is the requested memory size to allocate, and
        ``mem_size`` is the rounded memory size to be allocated, and
        ``mem_ptr`` is the memory pointer, and
        ``pmem_id`` is the pooled memory object ID.

    Attributes:
        file: Output file_like object that redirect to.
        flush: If ``True``, this hook forcibly flushes the text stream
            at the end of print. The default is ``True``.

    """

    name = 'DebugPrintHook'

    def __init__(self, file=sys.stdout, flush=True):
        self.file = file
        self.flush = flush

    def _print(self, msg):
        self.file.write(msg)
        self.file.write('\n')
        if self.flush:
            self.file.flush()

    def alloc_postprocess(self, **kwargs):
        msg = '{"hook":"%s","device_id":%d,' \
              '"mem_size":%d,"mem_ptr":%d}'
        msg %= ('alloc', kwargs['device_id'],
                kwargs['mem_size'], kwargs['mem_ptr'])
        self._print(msg)

    def malloc_postprocess(self, **kwargs):
        msg = '{"hook":"%s","device_id":%d,"size":%d,' \
              '"mem_size":%d,"mem_ptr":%d,"pmem_id":"%s"}'
        msg %= ('malloc', kwargs['device_id'], kwargs['size'],
                kwargs['mem_size'], kwargs['mem_ptr'], hex(kwargs['pmem_id']))
        self._print(msg)

    def free_postprocess(self, **kwargs):
        msg = '{"hook":"%s","device_id":%d,' \
              '"mem_size":%d,"mem_ptr":%d,"pmem_id":"%s"}'
        msg %= ('free', kwargs['device_id'],
                kwargs['mem_size'], kwargs['mem_ptr'], hex(kwargs['pmem_id']))
        self._print(msg)