import click

from sqlalchemy import insert

from wurker.settings import conf
from wurker.db import get_drone_status
from wurker.cli.base import BaseCommand
from wurker.drone import make_drone_clients
from wurker.models import WurkerBee, WurkerCron, WurkerQueue, WurkerCommand


class AddCommand(BaseCommand):
    """Encapsulate add command."""

    def __init__(self, logger=None):
        BaseCommand.__init__(self, logger=logger, check_schema=True)

        # We might need drone API auth
        self.drone_api_auth = conf.ApiAuth


    def run(self, option:str, data:dict, drones:list=[]):
        """Add a Wurker definition."""

        # Run via Drone API?
        if drones:
            return self.run_drone(option, add, drones)

        if not data:
            raise ValueError(f"You must set JSON for the {option} to add")

        # Make a copy of data so we don't change the input dict
        add = data.copy()

        if option == "command":
            return self.add_command(add)
        elif option == "cron":
            return self.add_cron(add)
        elif option == "bee":
            return self.add_bee(add)
        elif option == "ondemand":
            return self.add_ondemand(add)
        else:
            raise ValueError(f"Invalid 'option' value: '{option}'")


    def add_command(self, data:dict):
        """Add a Wurker Command definition."""

        if "app_id" not in data or not data["app_id"]:
            data["app_id"] = self.app_id

        # TODO Validate data more
        allowed_keys = {col.name for col in WurkerCommand.__table__.columns}
        invalid_keys = set(data.keys()) - allowed_keys
        if invalid_keys:
            raise ValueError(f"Invalid fields provided: {invalid_keys}")

        with self.RemoteSession() as db:
            query = db.query(WurkerCommand).filter(
                WurkerCommand.app_id == self.app_id,
                WurkerCommand.slug == data["slug"],
            )

            if query.count():
                raise ValueError(f"Wurker Command already exists with slug: '{data['slug']}'")

            stmt = insert(WurkerCommand).values(**data)
            db.execute(stmt)
            db.commit()


    def add_cron(self, data:dict):
        """Add a Wurker Cron definition."""

        if "app_id" not in data or not data["app_id"]:
            data["app_id"] = self.app_id

        # TODO Validate data more
        allowed_keys = {col.name for col in WurkerCron.__table__.columns}
        invalid_keys = set(data.keys()) - allowed_keys
        if invalid_keys:
            raise ValueError(f"Invalid fields provided: {invalid_keys}")

        with self.RemoteSession() as db:
            query = db.query(WurkerCron).filter(
                WurkerCron.app_id == self.app_id,
                WurkerCron.slug == data["slug"],
            )

            if query.count():
                raise ValueError(f"Wurker Cron already exists with slug: '{data['slug']}'")

            stmt = insert(WurkerCron).values(**data)
            db.execute(stmt)
            db.commit()


    def add_bee(self, data:dict):
        """Add a Wurker Bee definition, with extra 'cron_slug' mapping."""

        if "app_id" not in data or not data["app_id"]:
            data["app_id"] = self.app_id

        # TODO Validate data more
        allowed_keys = {col.name for col in WurkerBee.__table__.columns}
        allowed_keys.add("command_slug")
        allowed_keys.add("cron_slug")

        invalid_keys = set(data.keys()) - allowed_keys
        if invalid_keys:
            raise ValueError(f"Invalid fields provided: {invalid_keys}")

        with self.RemoteSession() as db:
            query = db.query(WurkerBee).filter(
                WurkerBee.app_id == self.app_id,
                WurkerBee.slug == data["slug"],
            )

            if query.count():
                raise ValueError(f"Wurker Bee already exists with slug: '{data['slug']}'")

            # Map to an existing Command at the same time using 'command_slug' feature?
            command = None
            if "command_slug" in data:
                command = db.query(WurkerCommand).filter(
                    WurkerCommand.app_id == self.app_id,
                    WurkerCommand.slug == data["command_slug"]
                ).first()

                if command is None:
                    raise ValueError(
                        f"Cannot set Wurker Command '{data['command_slug']}' for Bee '{data['slug']}': " +
                        "No Command record found"
                    )

                # We have the Command, so delete the 'command_slug' pseudo-field
                del data["command_slug"]
                data["command_id"] = command.id

            # Map to an existing Cron at the same time using 'cron_slug' feature?
            cron = None
            if "cron_slug" in data:
                cron = db.query(WurkerCron).filter(
                    WurkerCron.app_id == self.app_id,
                    WurkerCron.slug == data["cron_slug"]
                ).first()

                if cron is None:
                    raise ValueError(
                        f"Cannot set Wurker Cron '{data['cron_slug']}' for Bee '{data['slug']}': No Cron record found"
                    )

                # We have the Cron, so delete the 'cron_slug' pseudo-field
                del data["cron_slug"]

            # ADD the Bee
            stmt = insert(WurkerBee).values(**data)
            db.execute(stmt)
            db.commit()

            # Map the Bee to a Cron?
            if cron is not None:
                bee = db.query(WurkerBee).filter(
                    WurkerBee.app_id == self.app_id,
                    WurkerBee.slug == data['slug']
                ).one()

                if bee is None:
                    raise ValueError(f"Cannot assign Wurker Cron to Bee '{data['slug']}': No Bee record found")

                bee.crons.clear()
                bee.crons.append(cron)
                db.commit()


    def add_ondemand(self, data:dict):
        """Add a Wurker on-demand Bee definition."""

        # TODO Validate data

        with self.RemoteSession() as db:
            for item in data:

                # TODO Ondemand override args are on you, bih
                bee_slug = None
                ondemand_args = None

                if isinstance(item, dict):
                    bee_slug = list(item.keys()).pop()
                    ondemand_args = item[bee_slug]
                elif isinstance(item, str):
                    bee_slug = item
                else:
                    raise ValueError(f"Invalid Wurker Bee definition: {item}")

                if not bee_slug:
                    raise ValueError(f"Wurker Bee not found: '{bee_slug}'")

                # Try to get the Bee
                bee = db.query(WurkerBee).filter(
                    WurkerBee.app_id == self.app_id,
                    WurkerBee.slug == bee_slug,
                ).first()

                if bee is None:
                    raise ValueError(f"Bee not found: '{bee_slug}'")

                # Prohibit double-booking
                queued = db.query(WurkerQueue).filter(
                    WurkerQueue.app_id == self.app_id,
                    WurkerQueue.wurker_bee_id == bee.id,
                ).count()

                if queued:
                    raise ValueError(f"On-demand Wurker Bee '{bee_slug}' is already queued")

                # User can override args this way
                record = WurkerQueue(
                    app_id = self.app_id,
                    wurker_bee_id = bee.id,
                    args = ondemand_args,
                )
                db.add(record)

            db.commit()


    def run_drone(self, option:str, data:dict, drones:list):
        """Run via Drone API(s)."""

        drone_status = get_drone_status(self.remote_db_conf, self.app_id, drones=drones)
        if drone_status is None or not len(drone_status):
            click.echo(f"WARNING: Queen ({self.app_id}) Drones have no 'wurker_status' data", err=True)
            return

        drone_clients = make_drone_clients(drone_status, auth=self.drone_api_auth)

        for app_id, drone in drone_clients.items():
            drone.post_add(option, data=data)
