import os
import psutil
import logging
import traceback

from datetime import datetime, timezone
from sqlalchemy.sql import delete, insert, select, update

from wurker.settings import conf
from wurker.models import WurkerHive
from wurker.db import make_engine
from wurker.util import init_logger


class Hive:
    """Class to encapsulate and provide convenience for accessing local DB."""

    logger = None
    engine = None
    hive = None

    def __init__(self, app_root=None, clear=False, logger=None):
        db_conf = conf.LocalDB
        db_file = os.path.join(app_root, db_conf.database) if app_root else os.path.abspath(db_conf.database)

        self.logger = logger or init_logger()

        self.log(f"Connecting to local DB: '{db_file}'...", prefix="Hive.__init__")

        self.engine = make_engine(db_conf=db_conf)

        if clear:
            self.clear()

        self.sync()

    def log(self, msg, exception=False, level=logging.INFO, prefix=''):
        """Adds Bee-specific metadata to log entries."""

        if exception:
            level = logging.ERROR
            msg += "\n{0}".format(traceback.format_exc())

        prefix = prefix or "Hive"
        self.logger.log(msg, level=level, prefix=prefix)

    def clear(self):
        """Clear the active hive DB."""

        try:
            with self.engine.connect() as db:
                stmt = WurkerHive.__table__.delete()
                db.execute(stmt)
        except Exception as e:
            self.log(f"{e}", level=logging.ERROR)

    def load(self):
        """Get bees currently active in the hive DB."""

        self.hive = {}
        self.log("Loading Wurker Hive from local DB", level=logging.DEBUG)

        with self.engine.connect() as db:
            stmt = select(WurkerHive)
            hive = db.execute(stmt).all()
            for row in hive:
                if row.wurker_cron_id not in self.hive:
                    self.hive[row.wurker_cron_id] = {}
                self.hive[row.wurker_cron_id][row.wurker_bee_id] = {
                    "pid": row.pid,
                    "timestamp": row.timestamp,
                    "runcount": row.runcount,
                }

    def get(self, bee_id, cron_id):
        """Get a bee from the hive by unique key."""

        with self.engine.connect() as db:
            stmt = select(WurkerHive).where(
                WurkerHive.wurker_bee_id == bee_id,
                WurkerHive.wurker_cron_id == cron_id
            )
            return db.execute(stmt).first()

    def add(self, bee:object):
        """Add a bee to the hive."""

        data = {
            "wurker_bee_id": bee.id,
            "wurker_cron_id": bee.cron_id,
            "pid": bee.pid,
            "runcount": bee.runcount,
            "timestamp": datetime.now(timezone.utc),
        }

        with self.engine.connect() as db:
            stmt = insert(WurkerHive).prefix_with(" OR IGNORE").values(**data)
            db.execute(stmt)

    def set_runcount(self, bee:object):
        """Update the runcount for the given bee for display."""

        with self.engine.connect() as db:
            stmt = update(WurkerHive).values(runcount=bee.runcount)
            db.execute(stmt)

    def delete(self, bee:object=None, pid:int=None):
        """Delete a bee from the hive."""

        pid = bee.pid if bee is not None else pid
        if pid is None:
            return

        with self.engine.connect() as db:
            stmt = delete(WurkerHive).where(
                WurkerHive.pid == pid
            )
            db.execute(stmt)

    def sync(self):
        """Deletes Bee PIDs that aren't actually running."""

        self.log("Syncing Wurker Hive with running processes", level=logging.DEBUG)

        self.load()
        for cron_id, bee in self.hive.items():
            for bee_id, item in bee.items():
                if not psutil.pid_exists(item["pid"]):
                    self.delete(pid=item["pid"])
