import argparse
from concurrent import futures
import contextlib
import itertools
import json
import logging
import sys
import time

import aruba_common


log = logging.getLogger(__name__)
logging.basicConfig(
    level=logging.DEBUG,
    filename='aruba.log',
    filemode='w',
    format='%(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

available_projects = aruba_common.cfg.get('cpus', {})
available_switches = aruba_common.cfg.get('switches', {})


def valid_project(project):
    return True if project in available_projects else False


@contextlib.contextmanager
def switch_session(switch_name):
    """Use this context manager to ensure that we logout of the switch session"""
    session = None
    try:
        session = SwitchSession(switch_name)
        yield session

    except RuntimeError:
        raise

    finally:
        session and session.end_session()


class SwitchSession:
    """
    Encapsulates all operations a switch session can perform.
    Prefer to use the switch_session context manager as it will call end_session for you.
    Otherwise make sure to call end_session (essentially a __del__() but we shouldn't
    rely on Python to run it, as if it fails to, we'll leave sessions active.
    """

    def __init__(self, switch_name):
        # If keep-alive needs to be disabled to improve performance ,we should do
        # that here by creating a Requests session object.
        self.cookie = aruba_common.login_os(switch_name)
        if self.cookie is None:
            raise RuntimeError('Switch login failed. Unable to get a cookie.')
        self.ip = aruba_common.cfg['switches'][switch_name]['ip']
        self.base_url = 'http://{}/rest/v{}/'.format(self.ip, aruba_common.api_version)
        self.switch_name = switch_name

    def end_session(self):
        """Logout of the session."""
        aruba_common.logout_os(self.base_url, self.cookie)

    def poe_bounce(self, projects, sleep_time=0):
        """
        Bounce power on specified project.
        :param str projects: projects to bounce
        :param int sleep_time: time (seconds) to wait before restoring power
        """
        if not isinstance(projects, list):
            projects = [projects]

        self.poe_disable(projects)
        if sleep_time:
            log.info('Waiting for {} seconds before restoring power.'.format(sleep_time))
            time.sleep(sleep_time)
        self.poe_enable(projects)

    def _api_worker(self, project, poe_state):
        """Worker code that runs on each thread"""
        if valid_project(project):
            operation = 'enable' if poe_state else 'disable'
            port_number = aruba_common.cfg['cpus'][project]['switch_port_{}'.format(self.switch_name[-1:])]
            log.info('Performing poe {} for project {} on port {}.'.format(operation, project, port_number))
            cmd = '/ports/{}/poe'.format(port_number)
            payload = json.dumps({"is_poe_enabled": poe_state})
            aruba_common.execute_api(self.base_url, cmd, self.cookie, payload)
        else:
            log.error('Invalid project {}, poe {} will not be attempted.'. format(project, poe_state))

    def poe_state_change(self, poe_state, projects):
        """
        Set poe state of projects
        :param bool poe_state: True to enable power, False to disable power
        :param projects: projects to be operated on
        """
        if not isinstance(projects, list):
            projects = [projects]

        with futures.ThreadPoolExecutor(3) as executor:
            executor.map(
                self._api_worker,
                projects,
                itertools.repeat(poe_state, len(projects))
            )

    def poe_disable(self, projects):
        """
        Disable power on specified project.
        :param str projects: projects to bounce
        """
        self.poe_state_change(False, projects)

    def poe_enable(self, projects):
        """
        Enable power on specified port.
        :param str projects: projects to bounce
        """
        self.poe_state_change(True, projects)

    def get_port_by_id(self, port_number):
        """
        Get info on specified port
        :param str port_number: port number
        :return: JSON response
        """
        port = '/ports/{}'.format(port_number)
        return aruba_common.execute_api_get(self.base_url, port, self.cookie)


def main():
    parser = argparse.ArgumentParser(description='Bounce, enable or disable a project.')
    # All arguments are positional and required
    parser.add_argument(
        'operation',
        default='bounce',
        choices=['bounce', 'enable', 'disable'],
        help='POE operation to perform'
    )

    # 'all' is only a valid option for cmd line
    available_projects['all'] = None
    parser.add_argument(
        'projects',
        default='fca',
        nargs='+',
        choices=available_projects,
        help='Project to operate on (default: %(default)s)'
    )
    parser.add_argument(
        'switch',
        choices=available_switches,
        help='Switch to operate on'
    )
    parser.add_argument(
        '-sleep',
        '--sleep',
        type=int,
        default=0,
        help='For bounce only, the time (seconds) between disabling and enabling the port'

    )
    args = parser.parse_args()

    if 'all' in args.projects:
        args.projects = available_projects

    with switch_session(args.switch) as session:
        if 'bounce' in args.operation:
            session.poe_bounce(args.projects, args.sleep)
        elif 'enable' in args.operation:
            session.poe_enable(args.projects)
        else:
            session.poe_disable(args.projects)

    sys.exit(0)


if __name__ == '__main__':
    main()