import click

from sqlalchemy import update

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, WurkerCommand, WurkerCron


class EditCommand(BaseCommand):
    """Encapsulate edit 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, key:str, data:dict, drones:list=[]):
        """Edit a Wurker definition."""

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

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

        edit = data.copy()

        if option == "command":
            return self._edit_command(key, edit)
        elif option == "cron":
            return self._edit_cron(key, edit)
        elif option == "bee":
            return self._edit_bee(key, edit)
        else:
            raise ValueError(f"Invalid 'option' value: '{option}'")


    def _edit_command(self, key:str, data:dict):
        """Edit a Wurker Command definition."""

        # 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:
            command = db.query(WurkerCommand).filter(
                WurkerCommand.app_id == self.app_id,
                WurkerCommand.slug == key
            ).first()

            if command is None:
                raise ValueError(f"Cannot edit Wurker Command '{key}': No record found")

            stmt = update(WurkerCommand).where(
                WurkerCommand.app_id == self.app_id,
                WurkerCommand.slug == key
            ).values(**data)

            db.execute(stmt)
            db.commit()


    def _edit_cron(self, key:str, data:dict):
        """Edit a Wurker Cron definition."""

        # 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:
            cron = db.query(WurkerCron).filter(
                WurkerCron.app_id == self.app_id,
                WurkerCron.slug == key
            ).first()

            if cron is None:
                raise ValueError(f"Cannot edit Wurker Cron '{key}': No record found")

            stmt = update(WurkerCron).where(
                WurkerCron.id == cron.id
            ).values(**data)

            db.execute(stmt)
            db.commit()


    def _edit_bee(self, key:str, data:dict):
        """Edit a Wurker Bee definition, with extra 'cron_slug' mapping."""

        # TODO Validate data more
        allowed_keys = {col.name for col in WurkerBee.__table__.columns}
        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:

            # Retrieve the Bee instance for this app and slug.
            bee = db.query(WurkerBee).filter(
                WurkerBee.app_id == self.app_id,
                WurkerBee.slug == key
            ).one()

            if bee is None:
                raise ValueError(f"Cannot edit Wurker Bee '{key}': No Bee record found")

            # TODO Remap this Bee to a different Cron (fix later with migration)?
            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 '{key}': No Cron record found"
                    )

                # Clear any existing mappings from the join table, add new mapping, and commit
                bee.crons.clear()
                bee.crons.append(cron)
                db.commit()

                # Clear the "cron" pseudo-field so we don't try to update w/it
                del data["cron_slug"]

            # Maybe we only changed Cron? Then we're done
            if not data:
                return

            # Proceed to the rest of the Bee definition update...
            stmt = update(WurkerBee).where(
                WurkerBee.app_id == self.app_id,
                WurkerBee.slug == key
            ).values(**data)

            db.execute(stmt)
            db.commit()


    def run_drone(self, option:str, key: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_edit(option=option, key=key, data=data)
