#!/usr/bin/env python3
import base64
import time
import functools
import urllib.parse

import requests
from ansible.errors import AnsibleActionFail
from ansible.plugins.action import ActionBase


class Aruba:
    def __init__(self, base_url):
        self._session = requests.session()
        self._base_url = base_url

    def _url(self, url):
        return urllib.parse.urljoin(self._base_url, url)

    def login(self, username, password):
        response = self._session.post(
            self._url("/rest/v4/login-sessions"),
            json={"userName": username, "password": password},
        )
        if response.status_code != requests.codes.created:
            raise AnsibleActionFail("Login failed")

    def logout(self):
        response = self._session.delete(self._url("/rest/v4/login-sessions"))
        if response.status_code != requests.codes.no_content:
            raise AnsibleActionFail("Logout failed")

    def restore(self, config):
        response = self._session.post(
            self._url(
                "/rest/v4/system/config/cfg_restore/payload"
            ),
            json={
                "config_base64_encoded": base64.b64encode(config.encode()),
                "is_forced_reboot_enabled": True,
            },
        )
        if response.status_code != requests.codes.accepted:
            raise AnsibleActionFail("Restore failed")

        response = self._session.get(
            self._url(
                "/rest/v4/system/config/cfg_restore/payload/status"
            )
        )
        print(response.text)


class ActionModule(ActionBase):

    _VALID_ARGS = frozenset(("username", "password", "config", "url"))

    def _require_arg(self, name):
        try:
            return self._task.args[name]
        except KeyError:
            raise AnsibleActionFail("Missing argument: {}".format(name))

    def run(self, tmp=None, task_vars=None):
        task_vars = task_vars or {}
        result = super().run(tmp, task_vars)

        base_url = self._task.args.get("url")
        username = self._require_arg("username")
        password = self._require_arg("password")
        config = self._require_arg("config")

        aruba = Aruba(base_url)
        aruba.login(username, password)

        try:
            aruba.restore(config)
        except:
            raise
        else:
            aruba.logout()

        return result