import logging.config
import time

import features

from app import app, secscan_api, prometheus
from workers.worker import Worker
from data.database import UseThenDisconnect
from data.model.tag import (get_tags_images_eligible_for_scan, get_tag_pk_field,
                            get_max_id_for_sec_scan, get_min_id_for_sec_scan)
from util.secscan.api import SecurityConfigValidator
from util.secscan.analyzer import LayerAnalyzer, PreemptedException
from util.migrate.allocator import yield_random_entries
from endpoints.v2 import v2_bp


BATCH_SIZE = 50
DEFAULT_INDEXING_INTERVAL = 30


logger = logging.getLogger(__name__)
unscanned_images_gauge = prometheus.create_gauge('unscanned_images',
                                                 'Number of images that clair needs to scan.')
max_unscanned_images_gauge = prometheus.create_gauge('max_unscanned_image_id',
                                                     'Max ID of the unscanned images.')

class SecurityWorker(Worker):
  def __init__(self):
    super(SecurityWorker, self).__init__()
    validator = SecurityConfigValidator(app.config)
    if validator.valid():
      self._target_version = app.config.get('SECURITY_SCANNER_ENGINE_VERSION_TARGET', 3)
      self._analyzer = LayerAnalyzer(app.config, secscan_api)

      # Get the ID of the first image we want to analyze.
      self._min_id = app.config.get('SECURITY_SCANNER_INDEXING_MIN_ID',
                                    get_min_id_for_sec_scan(self._target_version))

      interval = app.config.get('SECURITY_SCANNER_INDEXING_INTERVAL', DEFAULT_INDEXING_INTERVAL)
      self.add_operation(self._index_images, interval)
    else:
      logger.warning('Failed to validate security scan configuration')

  def _index_images(self):
    def batch_query():
      return get_tags_images_eligible_for_scan(self._target_version)

    # Get the ID of the last image we can analyze. Will be None if there are no images in the
    # database.
    max_id = get_max_id_for_sec_scan()
    if max_id is None:
      return

    max_unscanned_images_gauge.Set(max_id)

    with UseThenDisconnect(app.config):
      to_scan_generator = yield_random_entries(
        batch_query,
        get_tag_pk_field(),
        BATCH_SIZE,
        max_id,
        self._min_id,
      )
      for candidate, abt, num_remaining in to_scan_generator:
        try:
          self._analyzer.analyze_recursively(candidate.image)
        except PreemptedException:
          logger.info('Another worker pre-empted us for layer: %s', candidate.id)
          abt.set()

        unscanned_images_gauge.Set(num_remaining)

    # If we reach this point, we analyzed every images up to max_id, next time the worker runs,
    # we want to start from the next image.
    self._min_id = max_id + 1

if __name__ == '__main__':
  app.register_blueprint(v2_bp, url_prefix='/v2')

  if not features.SECURITY_SCANNER:
    logger.debug('Security scanner disabled; skipping SecurityWorker')
    while True:
      time.sleep(100000)

  logging.config.fileConfig('conf/logging_debug.conf', disable_existing_loggers=False)
  worker = SecurityWorker()
  worker.start()