import os
import time
from datetime import datetime
import psutil
from typing import Any
import multiprocessing as mp
import socket

from py10x_kernel import OsUser
from core_10x.traitable import Traitable, T, RT
from core_10x.xdate_time import XDateTime
from core_10x.environment_variables import EnvVars
from core_10x.resource import NULL_RESOURCE
from core_10x.ts_store import TsStore

class PerfTimer:
    def __enter__(self):
        self.start = time.perf_counter_ns()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end = time.perf_counter_ns()
        self.elapsed = self.end - self.start

class LogMessage(Traitable, custom_collection = True, keep_history = False):
    ns: int             = T(T.ID)
    level: int          = T()
    mem_pc: float       = T()
    num_threads: int    = T()
    payload: Any        = T()

def logger_process(
        queue: mp.Queue,
        coll_name: str,
        do_print: bool
):
    uri = EnvVars.log_ts_store_uri
    log_store = TsStore.instance_from_uri(uri) if uri else NULL_RESOURCE

    data = queue.get()
    if data is None:
        return

    with log_store:
        first_msg = LogMessage(_replace = True, _collection_name = coll_name, **data)
        #start: datetime = first_msg.payload
        t1 = first_msg.ns
        first_msg.save()

    while True:
        data = queue.get()
        if data is None:
            break

        with log_store:
            msg = LogMessage(_replace = True, _collection_name = coll_name, **data)
            msg.save()

        if do_print:
            t2 = msg.ns
            dt = t2 - t1
            t1 = t2
            print(f'{(dt) / 1.e9}: {msg.payload}')

class Logger:
    def __init__(
        self,
        app_name: str,
        log_level: int,
        started_at: datetime = None,
        do_print = True
    ):
        self.app_name = app_name
        self.log_level = log_level
        if started_at is None:
            started_at = datetime.utcnow()
        self.started_at = started_at
        pid = os.getpid()
        full_name = f'{app_name}/{OsUser.me.name()}/{socket.gethostname()}/{XDateTime.datetime_to_str(started_at, False)}/{pid}'
        self.ps = psutil.Process(pid)

        self.queue = mp.Queue()
        self.proc = mp.Process(target = logger_process, args = (self.queue, full_name, do_print))
        self.proc.daemon = True
        self.proc.start()

    def log(self, data: dict):
        self.queue.put(data)

    def shutdown(self):
        self.queue.put(None)
        self.proc.join()

LOGGER: Logger = None
class LOG:
    def _log(log_level: int, payload: Any):
        assert LOGGER, 'You must have called LOG.begin(app_name, log_level)'
        if log_level > LOGGER.log_level:
            return

        if payload is not None:
            data = dict(
                ns          = time.perf_counter_ns(),
                level       = log_level,
                mem_pc      = LOGGER.ps.memory_percent(),
                num_threads = LOGGER.ps.num_threads(),
                payload     = payload
            )
        else:
            data = None

        LOGGER.log(data)

    @staticmethod
    def begin(app_name: str, log_level):
        global LOGGER
        if LOGGER is None:
            LOGGER = Logger(app_name, log_level.value)
            LOG.BRIEF(datetime.utcnow())

    @staticmethod
    def end():
        LOG._log(0, None)

    def __new__(cls, payload: Any): cls._log(0, payload)

    def BRIEF(payload: Any):    LOG._log(0, payload)
    BRIEF.value = 0
    def MEDIUM(payload: Any):   LOG._log(1, payload)
    MEDIUM.value = 1
    def DETAILED(payload: Any): LOG._log(2, payload)
    DETAILED.value = 2
    def VERBOSE(payload: Any):  LOG._log(3, payload)
    VERBOSE.value = 3

class LogReader:
    ...

