#!/usr/bin/env python
# vim: ai ts=4 sts=4 et sw=4
#
# Copyright (C) 2010, 2011, 2012, 2013, 2014 Intel, Inc.
#
#    This program is free software; you can redistribute it and/or
#    modify it under the terms of the GNU General Public License
#    as published by the Free Software Foundation; version 2
#    of the License.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, write to the Free Software
#    Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#
"""
"""

import os
import sys
import base64
from datetime import datetime
import requests
from common.utils import sync, unicode_to_str
from common.buildtrigger import trigger_next
import xml.etree.ElementTree as ElementTree
import ConfigParser

OBS_NORMAL_WORKER = 'obsnw'
OBS_POWER_WORKER  = 'obspw'

def get_worker_status():

    curr_time = datetime.now().strftime('%Y%m%d.%H%M%S')
    print curr_time

    obs_cred = '%s:%s' % (os.getenv('OBS_API_USERNAME'), os.getenv('OBS_API_PASSWD'))
    obs_host = os.getenv('OBS_API_URL').replace('//', '//%s@' % obs_cred)
    url = os.path.join(obs_host, 'build', '_workerstatus')

    resp = requests.get(url)
    print resp.text

    tree = ElementTree.fromstring(resp.text)

    ret_data = {'idle_normal': 0, 'idle_power': 0,
                'building_normal': 0, 'building_power': 0, 'building_power_chromium': 0,
                'waiting': 0, 'blocked': 0,
                'normal_slots': [], 'power_slots': [],
                'normal_slots_idle': [], 'power_slots_idle': []}

    worker_stat_container = {}

    for item in tree.findall('idle'):
        slot_inst = item.get('workerid')
        slot_name = slot_inst.split(':')[0]
        slot_number = slot_inst.split(':')[0].split(OBS_NORMAL_WORKER)[-1].split(OBS_POWER_WORKER)[-1]
        inst_number = slot_inst.split(':')[1]
        if slot_name not in worker_stat_container:
            worker_stat_container[slot_name] = []
        if inst_number in worker_stat_container[slot_name]:
            continue
        worker_stat_container[slot_name].append(inst_number)
        if slot_inst.startswith(OBS_NORMAL_WORKER):
            ret_data['idle_normal'] += 1
            if slot_number not in ret_data['normal_slots']:
                ret_data['normal_slots'].append(slot_number)
            if slot_number not in ret_data['normal_slots_idle']:
                ret_data['normal_slots_idle'].append(slot_number)
        elif slot_inst.startswith(OBS_POWER_WORKER):
            ret_data['idle_power'] += 1
            if slot_number not in ret_data['power_slots']:
                ret_data['power_slots'].append(slot_number)
            if slot_number not in ret_data['power_slots_idle']:
                ret_data['power_slots_idle'].append(slot_number)

    for item in tree.findall('building'):
        slot_inst = item.get('workerid')
        slot_name = slot_inst.split(':')[0]
        slot_number = slot_inst.split(':')[0].split(OBS_NORMAL_WORKER)[-1].split(OBS_POWER_WORKER)[-1]
        inst_number = slot_inst.split(':')[1]
        if slot_name not in worker_stat_container:
            worker_stat_container[slot_name] = []
        if inst_number in worker_stat_container[slot_name]:
            continue
        worker_stat_container[slot_name].append(inst_number)
        if slot_inst.startswith(OBS_NORMAL_WORKER):
            ret_data['building_normal'] += 1
            if slot_number not in ret_data['normal_slots']:
                ret_data['normal_slots'].append(slot_number)
            if slot_number in ret_data['normal_slots_idle']:
                ret_data['normal_slots_idle'].remove(slot_number)
        elif slot_inst.startswith(OBS_POWER_WORKER):
            ret_data['building_power'] += 1
            if slot_number not in ret_data['power_slots']:
                ret_data['power_slots'].append(slot_number)
            if item.get('package') == 'chromium-efl':
                ret_data['building_power_chromium'] += 1
            if slot_number in ret_data['power_slots_idle']:
                ret_data['power_slots_idle'].remove(slot_number)

    for item in tree.findall('waiting'):
        ret_data['waiting'] += int(item.get('jobs'))
    for item in tree.findall('blocked'):
        ret_data['blocked'] += int(item.get('jobs'))

    print '\n'
    for rd in ret_data:
        print rd, ret_data[rd]
    print '\n'
    sys.stdout.flush()

    if True:
        sync_src = os.path.join(os.getenv('JENKINS_HOME'), '.obs_worker_trend', 'obs_worker_history.log')
        sync_dest = os.path.join(os.getenv('IMG_SYNC_DEST_BASE'), 'snapshots', \
                                 '.dashboard', 'obs_worker_trend')
        curr_item = '%s,%s,%s,%s,%s,%s,%s\n' \
                     % (curr_time,
                        ret_data['idle_normal'], ret_data['idle_power'],
                        ret_data['building_normal'], ret_data['building_power'],
                        ret_data['waiting'], ret_data['blocked'])
        last_item = ''
        with open(sync_src, 'rb') as rh:
            rh.seek(-1024, os.SEEK_END)
            last_item = rh.readlines()[-1].decode()
        # Do not log if the last is the same as current
        if unicode_to_str(curr_item).split(',')[1:-1] != unicode_to_str(last_item).split(',')[1:-1]:
            with open(sync_src, 'a') as wh:
                wh.write(curr_item)
            print sync(os.path.dirname(sync_src), sync_dest)
            print curr_item
    return ret_data

def request_workers(num_executors, worker_type=OBS_NORMAL_WORKER):
    # Request number of imager nodes
    if num_executors <= 0:
        num_executors = 1
    print 'Requesting %d executors' % num_executors
    purpose = "OBS_WORKER_NORMAL"
    if worker_type == OBS_POWER_WORKER:
        purpose = "OBS_WORKER_POWER"
    if os.getenv("ONDEMAND_SLAVE_CONFIGURATION_ENABLED", "0") == "1":
        if num_executors > 0:
            trigger_next("SLAVE_BUILDER_%s" % worker_type, {"data":"dummy"}, \
                     extra_params={"ACTION": "REQUEST_WORKER", \
                                   "PURPOSE": purpose, \
                                   "REQUESTED_NUM_EXECUTORS": "%d" % num_executors})

def revoke_workers(slot_numbers, worker_type=OBS_NORMAL_WORKER):
    # Request number of imager nodes
    print 'Revoking %s - %s slots' % (worker_type, slot_numbers)
    purpose = "OBS_WORKER_NORMAL"
    if worker_type == OBS_POWER_WORKER:
        purpose = "OBS_WORKER_POWER"
    if os.getenv("ONDEMAND_SLAVE_CONFIGURATION_ENABLED", "0") == "1":
        if len(slot_numbers) > 0:
            trigger_next("SLAVE_BUILDER", {"data":"dummy"}, \
                     extra_params={"ACTION": "REVOKE_WORKER", \
                                   "PURPOSE": purpose, \
                                   "SLOT_NUMBERS": "%s" % ','.join(slot_numbers)})

class ReadConfig(object):

    def __init__(self, fname=None):
        if fname is None:
            conf_file = os.path.join(os.getenv('JENKINS_HOME'), 'init.groovy.d', 'setup.properties')
        else:
            conf_file = fname

        if not os.path.isfile(conf_file):
            return

        self.configParser = ConfigParser.RawConfigParser()
        with open(conf_file, 'r') as rf:
            temp_conf_str = '[default]\n' + rf.read()
        temp_conf_file = os.path.join(os.getenv('WORKSPACE'), '.prop')
        with open(temp_conf_file, 'w') as wf:
            wf.write(temp_conf_str.replace('\\n\\',''))

        self.configParser.read(temp_conf_file)

    def get_config(self, key, section='default'):
        return self.configParser.get(section, key)

def main():

    worker_status = get_worker_status()

    need_new_worker = worker_status['waiting'] - worker_status['idle_normal'] - worker_status['idle_power']

    # number of packages can be handled with existing instances
    num_can_handle_now = int(1 * (worker_status['idle_normal'] )) \
                       + int(1 * (worker_status['idle_power'] ))

    # number of packages need workers
    num_need_more = int((need_new_worker - num_can_handle_now))

    # If all the power workers are working for chromium-efl
    #  and there exist another waiting queue, 
    #  forcely create another power worker instance regardless of scheduling scheme.
    if ( num_need_more <= 0 \
         and worker_status['waiting'] > 0 \
         and worker_status['idle_power'] == 0 \
         and worker_status['building_power'] == worker_status['building_power_chromium'] ):
        num_need_more = 1

    print 'need_new_worker: %d' % need_new_worker
    print 'num_can_handle_now: %d' % num_can_handle_now
    print 'num_need_more: %d' % num_need_more

    if os.getenv('OBS_WORKER_NORMAL_AUTO_SCAILING_ENABLED', '0') == '0':
        return

    conf_inst = ReadConfig()
    max_normal = conf_inst.get_config('EC2_WORKER_OBS_NORMAL_INSTANCE_CAP_STR')
    max_power  = conf_inst.get_config('EC2_WORKER_OBS_POWER_INSTANCE_CAP_STR')

    if num_need_more > 0:
        if len(worker_status['power_slots']) < int(max_power):
            num_need_more = min(num_need_more / 4, 10)
            request_workers(num_need_more, worker_type=OBS_POWER_WORKER)
            request_workers(1, worker_type=OBS_NORMAL_WORKER)
            print "\"TitleDisplay\": \"+P(%d)\"" % num_need_more
        else:
            request_workers(num_need_more / 10, worker_type=OBS_NORMAL_WORKER)
            print "\"TitleDisplay\": \"+N(%d)\"" % num_need_more

    elif worker_status['waiting'] <= 0 and len(worker_status['normal_slots_idle']) > 0:
        revoke_workers(sorted(worker_status['normal_slots_idle']), OBS_NORMAL_WORKER)
        print "\"TitleDisplay\": \"-N(%d)\"" % len(worker_status['normal_slots_idle'])

    elif worker_status['waiting'] <= 0 and len(worker_status['power_slots_idle']) > 0:
        revoke_workers(sorted(worker_status['power_slots_idle']), OBS_POWER_WORKER)
        print "\"TitleDisplay\": \"-P(%d)\"" % len(worker_status['power_slots_idle'])

    else:
        print "\"TitleDisplay\": \"BN(%d) BP(%d) B(%d)\"" \
            % (worker_status['building_normal'], worker_status['building_power'], worker_status['blocked'])

if __name__ == '__main__':
    sys.exit(main())

