import tracemalloc
import json
import time
import logging
import os
import functools
log = logging.getLogger(__name__)


def _nested_tuple_from_nested_list(nl):
    if isinstance(nl, list):
        tuple_base = (_nested_tuple_from_nested_list(nested) for nested in nl)
        return tuple(tuple_base)
    return nl


class TracemallocJsonEncoder(json.JSONEncoder):
    def default(self, o):
        if isinstance(o, tracemalloc.Snapshot):
            return o.__dict__
        elif isinstance(o, tracemalloc._Traces):
            return list(o)
        elif isinstance(o, tracemalloc.Trace):
            return o._trace
        return super().default(o)


class TracemallocJsonDecoder(json.JSONDecoder):
    """
    Deserializes tracemalloc.Snapshot instances created with TracemallocJsonEncoder
    use with json.load/loads - pass class (not instance) in cls argument
    """
    def decode(self, s):
        raw = super().decode(s)
        return tracemalloc.Snapshot(
            traces=[_nested_tuple_from_nested_list(_t) for _t in raw['traces']],
            traceback_limit=raw['traceback_limit']
        )


class TracemallocWorker:
    def __init__(self):
        pass

    def set_tracking(self, tracemalloc_depth):
        if tracemalloc_depth > 0 and tracemalloc.is_tracing() and tracemalloc_depth != tracemalloc.get_traceback_limit():
            tracemalloc.stop()
            tracemalloc.start(tracemalloc_depth)
            pass

        if tracemalloc_depth > 0 and tracemalloc.is_tracing() == False:
            tracemalloc.start(tracemalloc_depth)

        if tracemalloc_depth == 0:
            tracemalloc.stop()

    def take_snapshot(self):
        if tracemalloc.is_tracing():
            return tracemalloc.take_snapshot()

    def write_snapshot(self, snapshot, dump_filename):
        try:
            with open(dump_filename, mode='w') as f:
                json.dump(snapshot, f, cls=TracemallocJsonEncoder)
        except:
            log.info("Unable to write tracemalloc json dump path: %s", dump_filename, exc_info=1)


class TracemallocDriver:
    def __init__(self, external_api):
        self.external_api = external_api
        self.snapshot_storage_path = self.external_api.get_log_path()
        log.info("Memory snapshot storage_path %s", self.snapshot_storage_path)
        self.worker = TracemallocWorker()

    def process_memory_usage(self):
        try:
            tracemalloc_depth = self.external_api.get_int_debug_flag("debugPluginAgentTracemallocDepthNative", 0)
            tracemalloc_dumps = self.external_api.get_bool_debug_flag("debugPluginAgentTracemallocDumpsNative", False)
            self.worker.set_tracking(tracemalloc_depth)

            if tracemalloc_depth > 0 and tracemalloc_dumps:
                log.info("Started taking memory snapshot")
                snapshot = self.worker.take_snapshot()
                if tracemalloc_dumps:
                    self.worker.write_snapshot(
                        snapshot,
                        dump_filename=self._snapshot_filename()
                    )
                log.info("Finished taking memory snapshot")
        except:
            log.info("Processing memory usage failed with exception", exc_info=1)

    def _snapshot_filename(self):
        pid = os.getpid()
        timestamp = str(time.monotonic()).replace('.', "_")
        return os.path.join(self.snapshot_storage_path, "pluginagent_tracemalloc_jsond_%s_%s.json" % (pid, timestamp))


def run_on_interval_decorator(f, interval):
    last_run = None
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        nonlocal last_run
        current_timestamp = time.monotonic()
        if not last_run or current_timestamp - last_run >= interval:
            last_run = current_timestamp
            f(*args, **kwargs)
    return wrapper

