import logging
import tracemalloc
from datetime import datetime
from tracemalloc import _format_size
from typing import Optional

from twisted.internet import reactor
from twisted.internet.threads import deferToThread

from peek_plugin_base.util.PeekPsUtil import PeekPsUtil

logger = logging.getLogger(__name__)


def rpad(val, count):
    val = str(val)
    return val + max(0, (count - len(val))) * " "


def lpad(val, count):
    val = str(val)
    return max(0, (count - len(val))) * " " + val


def center(val, width=80):
    return max(0, int((width - len(val)) / 2)) * " " + val


PEEK_MEM_DUMP_STACKTRACE = 2**0  # 1
PEEK_MEM_DUMP_VORTEX_OBSERVABLE_CACHE = 2**2  # 2
PEEK_MEM_DUMP_VORTEX_JSONABLE = 2**3  # 4
PEEK_MEM_DUMP_VORTEX_PUSH_PRODUCER = 2**4  # 8


def _formatTracemallocTraceback(top, size, count):
    sitePkgs = "site-packages"

    snapshot = tracemalloc.take_snapshot()

    # Get the objects by their line of malloc
    tracebackStats = list(
        filter(
            lambda s: s.size > size or s.count > count,
            snapshot.statistics("traceback"),
        )
    )[:top]

    text = "The traceback module is using %s memory\n" % _format_size(
        tracemalloc.get_tracemalloc_memory(), False
    )

    if not tracebackStats:
        text += "There are no large tracemalloc objects from traceback\n"
        return text

    text += "\n"

    # ---------------
    # Print the traceback summary

    for stat in tracebackStats:
        average = stat.size / stat.count

        tracePath = str(stat.traceback)
        if sitePkgs in tracePath:
            tracePath = tracePath[
                tracePath.index(sitePkgs) + len(sitePkgs) + 1 :
            ]

        text += "   size:%s" % lpad(_format_size(stat.size, False), 10)
        text += "   count:%s" % lpad(stat.count, 10)
        text += "   average:%s" % lpad(_format_size(average, False), 10)
        text += "   " + tracePath

        text += "\n"

        for line in stat.traceback.format():
            text += (" " * 8) + line + "\n"

        text += "\n"

    # ---------------
    # Print the lineno summary

    text += (
        " "
        + rpad("COUNT", 10)
        + " "
        + rpad("SIZE", 10)
        + " "
        + rpad("AVERAGE", 10)
        + " "
        + "MALLOC LINE"
        + "\n"
    )

    for stat in tracebackStats:
        average = stat.size / stat.count

        tracePath = str(stat.traceback)
        if sitePkgs in tracePath:
            tracePath = tracePath[
                tracePath.index(sitePkgs) + len(sitePkgs) + 1 :
            ]

        text += (
            " "
            + rpad(str(stat.count), 10)
            + " "
            + rpad(_format_size(stat.size, False), 10)
            + " "
            + rpad(_format_size(average, False), 10)
            + " "
            + tracePath
            + "\n"
        )

    text += "\n"

    return text


def _formatJsonableSummary(top, size):
    from vortex.Jsonable import Jsonable

    jsonableDump = Jsonable.memoryLoggingDump(top=top, over=size)

    if not jsonableDump:
        return "There are no large Vortex Jsonable Objects\n"

    text = " " + rpad("COUNT", 10) + " " + "OBJECT TYPE" + "\n"
    for objectType, count in jsonableDump:
        text += " " + rpad(str(count), 10) + " " + objectType + "\n"

    return text


def _formatObservableCacheSummary(top, size):
    from vortex.handler.TupleDataObservableCache import _CachedSubscribedData

    vortexCacheDump = _CachedSubscribedData.memoryLoggingDump(
        top=top, over=size
    )

    if not vortexCacheDump:
        return "There are no large Vortex Observable Caches\n"

    text = (
        " "
        + rpad("QUEUE ITEMS", 10)
        + " "
        + rpad("TOTAL SIZE", 10)
        + " "
        + "REMOTE VORTEX"
        + "\n"
    )

    for objectType, count, total in vortexCacheDump:
        text += (
            " "
            + rpad(str(count), 10)
            + " "
            + rpad(str(total), 10)
            + " "
            + objectType
            + "\n"
        )

    return text


def _formatVortexPushProducerSummary(top, msgs):
    from vortex.VortexWritePushProducer import VortexWritePushProducer

    stats = VortexWritePushProducer.memoryLoggingDump(top=top, msgs=msgs)

    if not stats:
        return "There are no large producer queues\n"

    text = (
        " "
        + rpad("MSGS", 10)
        + " "
        + rpad("TOTAL", 10)
        + " "
        + "REMOTE VORTEX"
        + "\n"
    )

    for objectType, count, total in stats:
        text += (
            " "
            + rpad(str(count), 10)
            + " "
            + rpad(str(total), 10)
            + " "
            + objectType
            + "\n"
        )

    return text


def dumpMemObjectToFile(
    debugMask: int = 0x11111111, serviceName: Optional[str] = None
):
    TRACEMALLOC_STACK_SIZE = 6

    # Start tracemalloc logging
    if debugMask & PEEK_MEM_DUMP_STACKTRACE:
        tracemalloc.start(TRACEMALLOC_STACK_SIZE)

    _dumpMemObjectToFile(debugMask, serviceName)


def _dumpMemObjectToFile(debugMask, serviceName):
    import os
    import pytz

    TRACEMALLOC_STACK_SIZE = 6

    class _State:
        lastClearDate = datetime.now(pytz.utc)

    TOP = 10
    COUNT_MIN = 10000
    TOTAL_SIZE = 1 * 1024 * 1024
    INDIVIDUAL_SIZE = 10 * 1024

    startTime = datetime.now(pytz.utc)
    # This is useful for debugging c binding memory leaks
    # roots = objgraph.get_leaking_objects()
    # objgraph.show_refs(roots[:3], refcounts=True, filename='/Users/peek/stash/roots.png')
    flushRequired = (
        3600 < (datetime.now(pytz.utc) - _State.lastClearDate).seconds
    )
    if flushRequired:
        _State.lastClearDate = datetime.now(pytz.utc)
    homeDir = os.path.expanduser("~/memdump-%s.log" % serviceName)
    with open(homeDir, "a") as f:
        # Write the start datetime
        f.write("=" * 80 + "\n")
        f.write("START - " + str(startTime) + "\n")

        if flushRequired:
            f.write("THE HOURLY FLUSH OCCURRED\n")

        # Write the total process memory
        f.write("-" * 80 + "\n")
        f.write(
            "Total python processes memory usage: "
            + rpad(_format_size(PeekPsUtil().memory_info.rss, False), 10)
            + "\n"
        )

        if debugMask & PEEK_MEM_DUMP_STACKTRACE:
            if flushRequired:
                tracemalloc.stop()
                tracemalloc.start(TRACEMALLOC_STACK_SIZE)

            f.write("-" * 80 + "\n")
            f.write(center("Python Tracemalloc Information") + "\n")
            f.write(_formatTracemallocTraceback(TOP, TOTAL_SIZE, COUNT_MIN))

        if debugMask & PEEK_MEM_DUMP_VORTEX_JSONABLE:
            # Write the _loop of the Jsonable objects
            f.write("-" * 80 + "\n")
            f.write(center("Vortex Jsonable Objects") + "\n")
            f.write(_formatJsonableSummary(TOP, INDIVIDUAL_SIZE))

        if debugMask & PEEK_MEM_DUMP_VORTEX_OBSERVABLE_CACHE:
            # Write the _loop of the cached vortex payloads
            f.write("-" * 80 + "\n")
            f.write(center("Vortex Observable Caches") + "\n")
            f.write(_formatObservableCacheSummary(TOP, INDIVIDUAL_SIZE))

        if debugMask & PEEK_MEM_DUMP_VORTEX_PUSH_PRODUCER:
            # Write the _loop of the cached vortex payloads
            f.write("-" * 80 + "\n")
            f.write(center("Vortex Write Push Producer") + "\n")
            f.write(_formatVortexPushProducerSummary(10, 1))

        # Write the end date
        f.write("-" * 80 + "\n")
        f.write(
            "END - Time taken : %s seconds"
            % str(datetime.now(pytz.utc) - startTime)
            + "\n"
        )
        f.write("-" * 80 + "\n")


def setupMemoryDebugging(serviceName: Optional[str] = None, debugMask: int = 0):
    import pytz

    TRACEMALLOC_STACK_SIZE = 6

    class _State:
        lastClearDate = datetime.now(pytz.utc)

    # Start tracemalloc logging
    if debugMask & PEEK_MEM_DUMP_STACKTRACE:
        tracemalloc.start(TRACEMALLOC_STACK_SIZE)

    # Start JSonable logging
    from vortex.Jsonable import Jsonable

    if debugMask & PEEK_MEM_DUMP_VORTEX_JSONABLE:
        Jsonable.setupMemoryLogging()

    from vortex.handler.TupleDataObservableCache import _CachedSubscribedData

    if debugMask & PEEK_MEM_DUMP_VORTEX_OBSERVABLE_CACHE:
        # Start cached encoded payload logging
        _CachedSubscribedData.setupMemoryLogging()

    from vortex.VortexWritePushProducer import VortexWritePushProducer

    if debugMask & PEEK_MEM_DUMP_VORTEX_PUSH_PRODUCER:
        # Start cached encoded payload logging
        VortexWritePushProducer.setupMemoryLogging()

    logger.warning("Memory Logging is enabled.")

    def _loop():
        dumpMemObjectToFile(debugMask, serviceName)

        reactor.callLater(60, deferToThread, _loop)

    _loop()
