# -*- coding: utf-8 -*-
'''
Management of the windows update agent
======================================

.. versionadded:: 2014.7.0

Set windows updates to run by category. Default behavior is to install
all updates that do not require user interaction to complete.

Optionally set ``category`` to a category of your choice to only
install certain updates. Default is to set to install all available updates.

The following example will install all Security and Critical Updates,
and download but not install standard updates.

.. code-block:: yaml

    updates:
      win_update.installed:
        - categories:
          - 'Critical Updates'
          - 'Security Updates'
        - skips:
          - downloaded
      win_update.downloaded:
        - categories:
          - 'Updates'
        - skips:
          - downloaded

You can also specify a number of features about the update to have a
fine grain approach to specific types of updates. These are the following
features/states of updates available for configuring:

.. code-block:: text

    'UI' - User interaction required, skipped by default
    'downloaded' - Already downloaded, included by default
    'present' - Present on computer, skipped by default
    'installed' - Already installed, skipped by default
    'reboot' - Reboot required, included by default
    'hidden' - Skip updates that have been hidden, skipped by default
    'software' - Software updates, included by default
    'driver' - driver updates, included by default

The following example installs all driver updates that don't require a reboot:
.. code-block:: yaml
    gryffindor:
      win_update.installed:
        - skips:
          - driver: True
          - software: False
          - reboot: False

To just update your windows machine, add this your sls:

.. code-block:: yaml
    updates:
      win_update.installed
'''

# Import Python libs
from __future__ import absolute_import
import logging

# Import 3rd-party libs
# pylint: disable=import-error
from salt.ext.six.moves import range  # pylint: disable=redefined-builtin
try:
    import win32com.client
    import pythoncom
    HAS_DEPENDENCIES = True
except ImportError:
    HAS_DEPENDENCIES = False
# pylint: enable=import-error

# Import salt libs
import salt.utils

log = logging.getLogger(__name__)


def __virtual__():
    '''
    Only works on Windows systems
    '''
    if salt.utils.is_windows() and HAS_DEPENDENCIES:
        return True
    return False


def _gather_update_categories(updateCollection):
    '''
    this is a convenience method to gather what categories of updates are available in any update
    collection it is passed. Typically though, the download_collection.
    Some known categories:
        Updates
        Windows 7
        Critical Updates
        Security Updates
        Update Rollups
    '''
    categories = []
    for i in range(updateCollection.Count):
        update = updateCollection.Item(i)
        for j in range(update.Categories.Count):
            name = update.Categories.Item(j).Name
            if name not in categories:
                log.debug('found category: {0}'.format(name))
                categories.append(name)
    return categories


class PyWinUpdater(object):
    def __init__(self, categories=None, skipUI=True, skipDownloaded=False,
            skipInstalled=True, skipReboot=False, skipPresent=False,
            skipSoftwareUpdates=False, skipDriverUpdates=False, skipHidden=True):
        log.debug('CoInitializing the pycom system')
        pythoncom.CoInitialize()

        # pylint: disable=invalid-name
        self.skipUI = skipUI
        self.skipDownloaded = skipDownloaded
        self.skipInstalled = skipInstalled
        self.skipReboot = skipReboot
        self.skipPresent = skipPresent
        self.skipHidden = skipHidden

        self.skipSoftwareUpdates = skipSoftwareUpdates
        self.skipDriverUpdates = skipDriverUpdates
        self.categories = categories
        self.foundCategories = None
        # pylint: enable=invalid-name

        log.debug('dispatching update_session to keep the session object.')
        self.update_session = win32com.client.Dispatch('Microsoft.Update.Session')

        log.debug('update_session got. Now creating a win_searcher to seek out the updates')
        self.win_searcher = self.update_session.CreateUpdateSearcher()

        # list of updates that are applicable by current settings.
        self.download_collection = win32com.client.Dispatch('Microsoft.Update.UpdateColl')

        # list of updates to be installed.
        self.install_collection = win32com.client.Dispatch('Microsoft.Update.UpdateColl')

        # the object responsible for fetching the actual downloads.
        self.win_downloader = self.update_session.CreateUpdateDownloader()
        self.win_downloader.Updates = self.download_collection

        # the object responsible for the installing of the updates.
        self.win_installer = self.update_session.CreateUpdateInstaller()
        self.win_installer.Updates = self.install_collection

        # the results of the download process
        self.download_results = None

        # the results of the installation process
        self.install_results = None

    def Search(self, searchString):
        try:
            log.debug('beginning search of the passed string: {0}'.format(searchString))
            self.search_results = self.win_searcher.Search(searchString)
            log.debug('search completed successfully.')
        except Exception as exc:
            log.info('search for updates failed. {0}'.format(exc))
            return exc

        log.debug('parsing results. {0} updates were found.'.format(
            self.search_results.Updates.Count))
        try:
            for update in self.search_results.Updates:
                if update.InstallationBehavior.CanRequestUserInput:
                    log.debug('Skipped update {0}'.format(update))
                    continue
                for category in update.Categories:
                    if self.skipDownloaded and update.IsDownloaded:
                        continue
                    if self.categories is None or category.Name in self.categories:
                        self.download_collection.Add(update)
                        log.debug('added update {0}'.format(update))
            self.foundCategories = _gather_update_categories(self.download_collection)
            return True
        except Exception as exc:
            log.info('parsing updates failed. {0}'.format(exc))
            return exc

    def AutoSearch(self):
        search_string = ''
        searchParams = []
        if self.skipInstalled:
            searchParams.append('IsInstalled=0')
        else:
            searchParams.append('IsInstalled=1')

        if self.skipHidden:
            searchParams.append('IsHidden=0')
        else:
            searchParams.append('IsHidden=1')

        if self.skipReboot:
            searchParams.append('RebootRequired=0')
        else:
            searchParams.append('RebootRequired=1')

        if self.skipPresent:
            searchParams.append('IsPresent=0')
        else:
            searchParams.append('IsPresent=1')

        if len(searchParams) > 1:
            for i in searchParams:
                search_string += '{0} and '.format(i)
        else:
            search_string += '{0} and '.format(searchParams[1])

        if not self.skipSoftwareUpdates and not self.skipDriverUpdates:
            search_string += 'Type=\'Software\' or Type=\'Driver\''
        elif not self.skipSoftwareUpdates:
            search_string += 'Type=\'Software\''
        elif not self.skipDriverUpdates:
            search_string += 'Type=\'Driver\''
        else:
            return False
            # if there is no type, the is nothing to search.
        log.debug('generated search string: {0}'.format(search_string))
        return self.Search(search_string)

    def Download(self):
        try:
            if self.download_collection.Count != 0:
                self.download_results = self.win_downloader.Download()
            else:
                log.debug('Skipped downloading, all updates were already cached.')
            return True
        except Exception as exc:
            log.debug('failed in the downloading {0}.'.format(exc))
            return exc

    def Install(self):
        try:
            for update in self.search_results.Updates:
                if update.IsDownloaded:
                    self.install_collection.Add(update)
            log.debug('Updates prepared. beginning installation')
        except Exception as exc:
            log.info('Preparing install list failed: {0}'.format(exc))
            return exc

        if self.install_collection.Count != 0:
            log.debug('Install list created, about to install')
            updates = []
            try:
                self.install_results = self.win_installer.Install()
                log.info('Installation of updates complete')
                return True
            except Exception as exc:
                log.info('Installation failed: {0}'.format(exc))
                return exc
        else:
            log.info('no new updates.')
            return True

    def GetInstallationResults(self):