from __future__ import annotations

import os
import re
import gc
import time
import socket
import pprint
import inspect
import threading
import functools
from contextlib import contextmanager
from typing import Optional, Union

import psutil

try:
    from greenlet import greenlet
except ImportError:
    greenlet = None
try:
    import gevent
    from gevent.socket import socket as gevent_socket
except ImportError:
    gevent_socket = socket.socket
try:
    from gevent.monkey import is_anything_patched
except ImportError:

    def is_anything_patched():
        return False


def _dictobj_replace(obj):
    if isinstance(obj, (str, bytes)):
        s = str(obj[:10])
        return s + "..."
    elif isinstance(obj, bytearray):
        s = str(obj[:10])
        return f"{(s[:-2])}...{(s[-2:])}"
    else:
        return obj


def print_obj(obj, indent=0):
    tab = " " * indent
    if inspect.ismodule(obj):
        print(f"{tab}+ {repr(obj)}")
    else:
        print(f"{tab}+ <{type(obj).__name__} object at 0x{'{:x}'.format(id(obj))}>")
        if isinstance(obj, dict):
            obj = {k: _dictobj_replace(v) for k, v in obj.items()}
        s = "".join(pprint.pformat(obj, indent=indent + 2, compact=True, depth=1))
        print(f"{tab}|_{s}")


def printobj_backref(queue, max_depth=10):
    ignore = set()
    ignore.add(id(ignore))
    ignore.add(id(queue))
    queue = [(obj, 0) for obj in queue]
    ignore.add(id(queue))
    for tpl in queue:
        ignore.add(id(tpl))
    gc.collect()
    while queue:
        target, depth = queue.pop()
        ignore.add(id(target))
        if depth > max_depth:
            continue
        print_obj(target, indent=depth)
        if inspect.ismodule(target):
            continue

        refs = gc.get_referrers(target)
        ignore.add(id(refs))
        for ref in refs:
            if id(ref) in ignore:
                continue
            tpl = (ref, depth + 2)
            ignore.add(id(tpl))
            queue.append(tpl)


def gc_collect(*classes):
    """
    :param classes:
    :returns set:
    """
    ret = []
    if not classes:
        return ret
    # TODO: gevent monitoring says this blocks
    # find a better way?
    yield_to_eventloop = is_anything_patched()
    for ob in gc.get_objects():
        if yield_to_eventloop:
            gevent.sleep()
        try:
            if not isinstance(ob, classes):
                continue
        except ReferenceError:
            continue
        ret.append(ob)
    return ret


def find_processes_by_file(file_path):
    processes_with_file_open = []

    for proc in psutil.process_iter(["pid", "name", "open_files"]):
        try:
            open_files = proc.info["open_files"]
            if open_files:
                for open_file in open_files:
                    if open_file.path == file_path:
                        processes_with_file_open.append(proc.info)
                        break
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            continue

    return processes_with_file_open


def file_descriptors(
    proc: Optional[Union[int, psutil.Process]] = None
) -> dict[int, str]:
    """
    All file descriptors of a process

    :param proc: current by default
    :returns dict: mapping of file id to path
    """
    if proc is None:
        proc = os.getpid()

    try:
        if not isinstance(proc, psutil.Process):
            proc = psutil.Process(pid=proc)
        return {fd.fd: fd.path for fd in proc.open_files()}
    except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
        return {}


def greenlets():
    """
    All greenlets of this process

    :returns set:
    """
    if greenlet is None:
        return set()
    return set(gc_collect(greenlet))


def sockets():
    """
    All sockets of this process

    :returns set:
    """
    return set(gc_collect(gevent.socket.socket, socket.socket))


def resources(*classes):
    """
    All other resource of this process

    :returns set:
    """
    return set(gc_collect(*classes))


def threads(pid=None):
    """
    All threads of a process

    :param int pid: current by default
    :returns set:
    """
    if pid is None or pid == os.getpid():
        return set(threading.enumerate())
    else:
        process = psutil.Process(pid=pid)
        return {t.id for t in process.threads()}


def memory(pid=None):
    """
    Used memory of a process

    :param int pid: current by default
    :returns int: bytes
    """
    process = psutil.Process(pid=pid)
    return process.memory_info().rss


def file_processes(pattern: str) -> list[tuple[str, psutil.Process]]:
    r"""
    All processes that have a file opened matching the pattern

    Example: all open hdf5 files

    .. code-block:: python
        file_processes(r".+\.h5$")

    :param pattern: regex pattern
    :returns: list of filenames and processes
    """
    ret = []
    regex = re.compile(pattern)

    for proc in psutil.process_iter(["pid", "name"]):
        try:
            fds = file_descriptors(proc)
            for filename in fds.values():
                if regex.match(filename):
                    ret.append((filename, proc))
        except Exception as e:
            print(f"Unexpected error with process {proc.info['pid']}: {e}")
            continue

    return ret


def log_file_processes(logfunc, pattern):
    """
    All processes that have a file opened

    :param callable logfunc:
    :param str pattern: regex pattern
    """
    procs = file_processes(pattern)
    if not procs:
        return
    pid = os.getpid()
    msg = "\n### Files matching {} opened by {} process".format(
        repr(pattern), len(procs)
    )
    for fname, proc in procs:
        cmdline = repr(" ".join(proc.cmdline()))
        if proc.pid == pid:
            proc = "CURRENT {}".format(proc)
        msg += "\nFile {}:\n Opened by {}\n {}".format(repr(fname), proc, cmdline)
    msg += "\n###\n"
    logfunc(msg)


def matching_processes(pattern):
    """
    Find processes

    :param str pattern: regex pattern
    :yields Process:
    """
    for proc in psutil.process_iter():
        if re.match(pattern, " ".join(proc.cmdline())):
            yield proc


def terminate(pattern, wait=False):
    """
    Terminate a process

    :param str pattern: regex pattern
    :param bool force:
    """
    for proc in matching_processes(pattern):
        proc.terminate()
        if wait:
            proc.wait()


def kill(pattern, wait=False):
    """
    Terminate a process

    :param str pattern: regex pattern
    :param bool force:
    """
    for proc in matching_processes(pattern):
        proc.kill()
        if wait:
            proc.wait()


def log_diff(logfunc, difffunc, typ, diffargs=None, prefix=""):
    """
    :param callable logfunc:
    :param callable difffunc:
    :param str typ:
    :param tuple diffargs:
    :param str prefix:
    """
    new, diff = difffunc(*diffargs)
    if diff:
        if prefix:
            prefix = prefix.format(len(diff)) + ": "
        else:
            prefix = f"{len(diff)} {typ} difference "
        try:
            lst = sorted(diff)
        except TypeError:
            lst = list(diff)
        printobj_backref(lst)
        lst = list(map(str, lst))
        logfunc(f"{prefix}{lst}")
    return new


def _fd_diff(old):
    """New file descriptors

    :param dict old:
    :returns dict, set:
    """
    new = file_descriptors()
    diff = set(new.values()) - set(old.values())
    return new, diff


def _class_diff(old, classes):
    """New instance diff

    :param set old:
    :param tuple classes:
    :returns set, set:
    """
    new = set(gc_collect(*classes))
    return new, new - old


def _greenlet_diff(old):
    """
    New greenlets

    :param set old:
    :returns set, set:
    """
    new = greenlets()
    return new, new - old


def _socket_diff(old):
    """
    New sockets

    :param set old:
    :returns set, set:
    """
    new = sockets()
    return new, new - old


def _thread_diff(old):
    """
    New threads

    :param set old:
    :returns set, set:
    """
    new = threads()
    return new, new - old


def log_fd_diff(logfunc, old, prefix=""):
    """
    Print new file descriptors

    :param callable logfunc:
    :param dict old:
    :param str prefix:
    :returns dict:
    """
    return log_diff(logfunc, _fd_diff, "fds", diffargs=(old,), prefix=prefix)


def log_greenlet_diff(logfunc, old, prefix=""):
    """
    Print new greenlets

    :param callable logfunc:
    :param set old:
    :param str prefix:
    :returns set:
    """
    return log_diff(
        logfunc, _greenlet_diff, "greenlets", diffargs=(old,), prefix=prefix
    )


def log_class_diff(logfunc, old, classes, prefix=""):
    """Print new class instances

    :param callable logfunc:
    :param set old:
    :param str prefix:
    :returns set:
    """
    return log_diff(
        logfunc, _class_diff, str(classes), diffargs=(old, classes), prefix=prefix
    )


def log_socket_diff(logfunc, old, prefix=""):
    """
    Print new sockets

    :param callable logfunc:
    :param set old:
    :param str prefix:
    :returns set:
    """
    return log_diff(logfunc, _socket_diff, "sockets", diffargs=(old,), prefix=prefix)


def log_thread_diff(logfunc, old, prefix=""):
    """
    Print new threads

    :param callable logfunc:
    :param set old:
    :param str prefix:
    :returns set:
    """
    return log_diff(logfunc, _thread_diff, "threads", diffargs=(old,), prefix=prefix)


def _diff_exception(*args, **kwargs):
    raise AssertionError(*args, **kwargs)


def assert_fd_diff(old, prefix=""):
    """
    Raise exception when new file descriptors

    :param dict old:
    :param str prefix:
    :returns dict:
    """
    return log_fd_diff(_diff_exception, old, prefix=prefix)


def assert_greenlet_diff(old, prefix=""):
    """
    Raise exception when new greenlets

    :param set old:
    :param str prefix:
    :returns set:
    """
    return log_greenlet_diff(_diff_exception, old, prefix=prefix)


def assert_class_diff(old, classes, prefix=""):
    """
    Raise exception when new instances

    :param set old:
    :param str prefix:
    :returns set:
    """
    return log_class_diff(_diff_exception, old, classes, prefix=prefix)


def assert_socket_diff(old, prefix=""):
    """
    Raise exception when new sockets

    :param set old:
    :param str prefix:
    :returns set:
    """
    return log_socket_diff(_diff_exception, old, prefix=prefix)


def assert_thread_diff(old, prefix=""):
    """
    Raise exception when new threads

    :param set old:
    :param str prefix:
    :returns set:
    """
    return log_thread_diff(_diff_exception, old, prefix=prefix)


class ResourceMonitor:
    def __init__(self, *classes, greenlets=True, threads=True, fds=True, sockets=True):
        self._yield_to_eventloop = is_anything_patched()
        self.classes = classes
        if classes:
            self.others = set()
        else:
            self.others = None
        if greenlets:
            self.greenlets = set()
        else:
            self.greenlets = None
        if threads:
            self.threads = set()
        else:
            self.threads = None
        if fds:
            self.fds = set()
        else:
            self.fds = None
        if sockets:
            self.sockets = set()
        else:
            self.sockets = None

    def start(self):
        if self._yield_to_eventloop:
            gevent.get_hub()
        if self.greenlets is not None:
            self.greenlets = greenlets()
        if self.threads is not None:
            self.threads = threads()
        if self.fds is not None:
            self.fds = file_descriptors()
        if self.sockets is not None:
            self.sockets = sockets()
        if self.others is not None:
            self.others = resources(*self.classes)

    @contextmanager
    def check_leaks_context(self, msg=None):
        self.start()
        try:
            yield
        finally:
            self.check_leaks(msg=msg)

    def wait_gc_collect(self):
        time.sleep(0.1)
        while gc.collect():
            time.sleep(0.1)

    def check_leaks(self, msg=None):
        self.wait_gc_collect()
        if self.others is not None:
            assert_class_diff(self.others, self.classes, prefix=msg)
        if self.sockets is not None:
            assert_socket_diff(self.sockets, prefix=msg)
        if self.fds is not None:
            assert_fd_diff(self.fds, prefix=msg)
        if self.threads is not None:
            assert_thread_diff(self.threads, prefix=msg)
        if self.greenlets is not None:
            assert_greenlet_diff(self.greenlets, prefix=msg)


def check_resource_leaks(*classes, **kw):
    """Decorator for function resource checking"""

    def _check_resource_leaks(method):
        @functools.wraps(method)
        def inner(*args, **kwargs):
            mon = ResourceMonitor(*classes, **kw)
            with mon.check_leaks_context(msg=method.__name__):
                method(*args, **kwargs)

        return inner

    return _check_resource_leaks
