import os
import psutil
import logging
from pathlib import Path


logging.basicConfig(
    format='%(asctime)s %(levelname)-8s %(message)s',
    level=logging.INFO,
    datefmt='%m/%d %H:%M:%S')
logger = logging.getLogger(__name__)


def log_sys_info():
    cpu_cores = psutil.cpu_count()
    # cpu_freq = '{:.2f}'.format(psutil.cpu_freq().max / 1000) + 'GHz'
    cpu_percent = '{:.2f}'.format(psutil.cpu_percent()) + '%'
    mem = psutil.virtual_memory()
    mem_total = '{:.2f}'.format(mem.total / 1024 / 1024 / 1024) + 'GB'
    mem_used = '{:.2f}'.format(mem.used / 1024 / 1024 / 1024) + 'GB'
    mem_percent = '{:.2f}'.format(mem.percent) + '%'
    disk = psutil.disk_usage('.')
    disk_total = '{:.2f}'.format(disk.total / 1024 / 1024 / 1024) + 'GB'
    disk_used = '{:.2f}'.format(disk.used / 1024 / 1024 / 1024) + 'GB'
    disk_percent = '{:.2f}'.format(disk.percent) + '%'

    logger.info('======== SYSTEM INFO =========')
    logger.info(f'CPU: {cpu_cores} cores, {cpu_percent} used')
    logger.info(f'RAM: {mem_used} / {mem_total}, {mem_percent} used')
    logger.info(f'DISK: {disk_used} / {disk_total}, {disk_percent} used')
    logger.info('==============================')


class CHAT_DB:
    def __init__(self):
        self.prompts = {}
        self.results = {}
        self.index = 1
        self.lockfile = '.lock'

    def set(self, index, prompt=None, result=None):
        assert prompt or result
        if prompt:
            if index in self.prompts:
                raise ValueError('Prompt already exists')
            self.prompts[index] = prompt
            self.index += 1
        if result:
            self.results[index] = result

    def lock(self):
        if not os.path.exists(self.lockfile):
            Path(self.lockfile).touch(exist_ok=True)

    def unlock(self):
        if os.path.exists(self.lockfile):
            os.remove(self.lockfile)

    def islocked(self):
        return os.path.exists(self.lockfile)

    def clean(self):
        if len(self.prompts) > 100:
            self.prompts = dict(list(self.prompts.items())[-100:])
        k = list(set(self.prompts.keys()).intersection(set(self.results.keys())))  # keys to preserve
        self.prompts = {i: self.prompts[i] for i in k}
        self.results = {i: self.results[i] for i in k}
        log_sys_info()


db = CHAT_DB()