#!/usr/bin/python
# -*- coding: utf-8 -*-

# (c) 2018, Chris Houseknecht <@chouseknecht>
# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)

from __future__ import absolute_import, division, print_function


__metaclass__ = type


DOCUMENTATION = r'''

module: k8s_scale

short_description: Set a new size for a Deployment, ReplicaSet, Replication Controller, or Job.

author:
    - "Chris Houseknecht (@chouseknecht)"
    - "Fabian von Feilitzsch (@fabianvf)"

description:
  - Similar to the kubectl scale command. Use to set the number of replicas for a Deployment, ReplicaSet,
    or Replication Controller, or the parallelism attribute of a Job. Supports check mode.
  - C(wait) parameter is not supported for Jobs.

extends_documentation_fragment:
  - kubernetes.core.k8s_name_options
  - kubernetes.core.k8s_auth_options
  - kubernetes.core.k8s_resource_options
  - kubernetes.core.k8s_scale_options

options:
  label_selectors:
    description: List of label selectors to use to filter results.
    type: list
    elements: str
    version_added: 2.0.0
  continue_on_error:
    description:
    - Whether to continue on errors when multiple resources are defined.
    type: bool
    default: False
    version_added: 2.0.0

requirements:
    - "python >= 3.6"
    - "kubernetes >= 12.0.0"
    - "PyYAML >= 3.11"
'''

EXAMPLES = r'''
- name: Scale deployment up, and extend timeout
  kubernetes.core.k8s_scale:
    api_version: v1
    kind: Deployment
    name: elastic
    namespace: myproject
    replicas: 3
    wait_timeout: 60

- name: Scale deployment down when current replicas match
  kubernetes.core.k8s_scale:
    api_version: v1
    kind: Deployment
    name: elastic
    namespace: myproject
    current_replicas: 3
    replicas: 2

- name: Increase job parallelism
  kubernetes.core.k8s_scale:
    api_version: batch/v1
    kind: job
    name: pi-with-timeout
    namespace: testing
    replicas: 2

# Match object using local file or inline definition

- name: Scale deployment based on a file from the local filesystem
  kubernetes.core.k8s_scale:
    src: /myproject/elastic_deployment.yml
    replicas: 3
    wait: no

- name: Scale deployment based on a template output
  kubernetes.core.k8s_scale:
    resource_definition: "{{ lookup('template', '/myproject/elastic_deployment.yml') | from_yaml }}"
    replicas: 3
    wait: no

- name: Scale deployment based on a file from the Ansible controller filesystem
  kubernetes.core.k8s_scale:
    resource_definition: "{{ lookup('file', '/myproject/elastic_deployment.yml') | from_yaml }}"
    replicas: 3
    wait: no

- name: Scale deployment using label selectors (continue operation in case error occured on one resource)
  kubernetes.core.k8s_scale:
    replicas: 3
    kind: Deployment
    namespace: test
    label_selectors:
      - app=test
    continue_on_error: true
'''

RETURN = r'''
result:
  description:
  - If a change was made, will return the patched object, otherwise returns the existing object.
  returned: success
  type: complex
  contains:
     api_version:
       description: The versioned schema of this representation of an object.
       returned: success
       type: str
     kind:
       description: Represents the REST resource this object represents.
       returned: success
       type: str
     metadata:
       description: Standard object metadata. Includes name, namespace, annotations, labels, etc.
       returned: success
       type: complex
     spec:
       description: Specific attributes of the object. Will vary based on the I(api_version) and I(kind).
       returned: success
       type: complex
     status:
       description: Current status details for the object.
       returned: success
       type: complex
     duration:
       description: elapsed time of task in seconds
       returned: when C(wait) is true
       type: int
       sample: 48
'''

import copy

from ansible_collections.kubernetes.core.plugins.module_utils.ansiblemodule import AnsibleModule
from ansible_collections.kubernetes.core.plugins.module_utils.args_common import (
    AUTH_ARG_SPEC, RESOURCE_ARG_SPEC, NAME_ARG_SPEC)


SCALE_ARG_SPEC = {
    'replicas': {'type': 'int', 'required': True},
    'current_replicas': {'type': 'int'},
    'resource_version': {},
    'wait': {'type': 'bool', 'default': True},
    'wait_timeout': {'type': 'int', 'default': 20},
    'wait_sleep': {'type': 'int', 'default': 5},
}


def execute_module(module, k8s_ansible_mixin,):
    k8s_ansible_mixin.set_resource_definitions(module)

    definition = k8s_ansible_mixin.resource_definitions[0]

    name = definition['metadata']['name']
    namespace = definition['metadata'].get('namespace')
    api_version = definition['apiVersion']
    kind = definition['kind']
    current_replicas = module.params.get('current_replicas')
    replicas = module.params.get('replicas')
    resource_version = module.params.get('resource_version')

    label_selectors = module.params.get('label_selectors')
    if not label_selectors:
        label_selectors = []
    continue_on_error = module.params.get('continue_on_error')

    wait = module.params.get('wait')
    wait_time = module.params.get('wait_timeout')
    wait_sleep = module.params.get('wait_sleep')
    existing = None
    existing_count = None
    return_attributes = dict(result=dict())
    if module._diff:
        return_attributes['diff'] = dict()
    if wait:
        return_attributes['duration'] = 0

    resource = k8s_ansible_mixin.find_resource(kind, api_version, fail=True)

    from ansible_collections.kubernetes.core.plugins.module_utils.common import NotFoundError

    multiple_scale = False
    try:
        existing = resource.get(name=name, namespace=namespace, label_selector=','.join(label_selectors))
        if existing.kind.endswith('List'):
            existing_items = existing.items
            multiple_scale = len(existing_items) > 1
        else:
            existing_items = [existing]
    except NotFoundError as exc:
        module.fail_json(msg='Failed to retrieve requested object: {0}'.format(exc),
                         error=exc.value.get('status'))

    if multiple_scale:
        # when scaling multiple resource, the 'result' is changed to 'results' and is a list
        return_attributes = {'results': []}
    changed = False

    def _continue_or_fail(error):
        if multiple_scale and continue_on_error:
            if "errors" not in return_attributes:
                return_attributes['errors'] = []
            return_attributes['errors'].append({'error': error, 'failed': True})
        else:
            module.fail_json(msg=error, **return_attributes)

    def _continue_or_exit(warn):
        if multiple_scale:
            return_attributes['results'].append({'warning': warn, 'changed': False})
        else:
            module.exit_json(warning=warn, **return_attributes)

    for existing in existing_items:
        if module.params['kind'] == 'job':
            existing_count = existing.spec.parallelism
        elif hasattr(existing.spec, 'replicas'):
            existing_count = existing.spec.replicas

        if existing_count is None:
            error = 'Failed to retrieve the available count for object kind={0} name={1} namespace={2}.'.format(
                    existing.kind, existing.metadata.name, existing.metadata.namespace)
            _continue_or_fail(error)
            continue

        if resource_version and resource_version != existing.metadata.resourceVersion:
            warn = 'expected resource version {0} does not match with actual {1} for object kind={2} name={3} namespace={4}.'.format(
                   resource_version, existing.metadata.resourceVersion, existing.kind, existing.metadata.name, existing.metadata.namespace)
            _continue_or_exit(warn)
            continue

        if current_replicas is not None and existing_count != current_replicas:
            warn = 'current replicas {0} does not match with actual {1} for object kind={2} name={3} namespace={4}.'.format(
                   current_replicas, existing_count, existing.kind, existing.metadata.name, existing.metadata.namespace)
            _continue_or_exit(warn)
            continue

        if existing_count != replicas:
            if not module.check_mode:
                if module.params['kind'] == 'job':
                    existing.spec.parallelism = replicas
                    result = resource.patch(existing.to_dict()).to_dict()
                else:
                    result = scale(module, k8s_ansible_mixin, resource, existing, replicas, wait, wait_time, wait_sleep)
                    changed = changed or result['changed']
        else:
            name = existing.metadata.name
            namespace = existing.metadata.namespace
            existing = resource.get(name=name, namespace=namespace)
            result = {'changed': False, 'result': existing.to_dict()}
            if module._diff:
                result['diff'] = {}
            if wait:
                result['duration'] = 0
        # append result to the return attribute
        if multiple_scale:
            return_attributes['results'].append(result)
        else:
            module.exit_json(**result)

    module.exit_json(changed=changed, **return_attributes)


def argspec():
    args = copy.deepcopy(SCALE_ARG_SPEC)
    args.update(RESOURCE_ARG_SPEC)
    args.update(NAME_ARG_SPEC)
    args.update(AUTH_ARG_SPEC)
    args.update({'label_selectors': {'type': 'list', 'elements': 'str', 'default': []}})
    args.update(({'continue_on_error': {'type': 'bool', 'default': False}}))
    return args


def scale(module, k8s_ansible_mixin, resource, existing_object, replicas, wait, wait_time, wait_sleep):
    name = existing_object.metadata.name
    namespace = existing_object.metadata.namespace
    kind = existing_object.kind

    if not hasattr(resource, 'scale'):
        module.fail_json(
            msg="Cannot perform scale on resource of kind {0}".format(resource.kind)
        )

    scale_obj = {'kind': kind, 'metadata': {'name': name, 'namespace': namespace}, 'spec': {'replicas': replicas}}

    existing = resource.get(name=name, namespace=namespace)

    try:
        resource.scale.patch(body=scale_obj)
    except Exception as exc:
        module.fail_json(msg="Scale request failed: {0}".format(exc))

    k8s_obj = resource.get(name=name, namespace=namespace).to_dict()
    match, diffs = k8s_ansible_mixin.diff_objects(existing.to_dict(), k8s_obj)
    result = dict()
    result['result'] = k8s_obj
    result['changed'] = not match
    if module._diff:
        result['diff'] = diffs

    if wait:
        success, result['result'], result['duration'] = k8s_ansible_mixin.wait(resource, scale_obj, wait_sleep, wait_time)
        if not success:
            module.fail_json(msg="Resource scaling timed out", **result)
    return result


def main():
    mutually_exclusive = [
        ('resource_definition', 'src'),
    ]
    module = AnsibleModule(argument_spec=argspec(), mutually_exclusive=mutually_exclusive, supports_check_mode=True)
    from ansible_collections.kubernetes.core.plugins.module_utils.common import (
        K8sAnsibleMixin, get_api_client)
    k8s_ansible_mixin = K8sAnsibleMixin(module)
    k8s_ansible_mixin.client = get_api_client(module=module)
    execute_module(module, k8s_ansible_mixin)


if __name__ == '__main__':
    main()
