#!/usr/bin/python
# -*- coding: utf-8 -*-
# vim: set expandtab:tabstop=4:softtabstop=4:shiftwidth=4

import argparse
from datetime import datetime
from fcntl import lockf, LOCK_EX
import logging
from kazoo.client import KazooClient
import os
from subprocess import check_call, check_output, Popen, PIPE
import yaml

with open('/etc/yandex/environment.type', 'r') as fh:
    if fh.readline().strip() not in ['production', 'prestable']:
        exit(0)

zk_node_format = '/direct/versions_current/{app}/{fqdn}/dpkg'
usage = u'''Записывает версии приложений, установленных на этой машине, в zk-ноды '%s'.
"Версия приложения" = версия пакета, соответствующего приложению в apps.conf.
Примеры вызова:

    # просто записать версии в zk
    dt-write-versions-to-zk
    # записать версии в zk и сделать отметку -- пустой файл с mtime, взятым от /var/lib/dpkg/status на момент запуска
    dt-write-versions-to-zk --set-mark /var/lib/dt-write-versions-to-zk/dpkg-status.mark

''' % zk_node_format

parser = argparse.ArgumentParser(description=usage, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('--set-mark', metavar='mark_file_path')
args = parser.parse_args()
mark_file_path = None
if args.set_mark:
    mark_file_path = args.set_mark

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s')
logging.info('START')

fqdn = check_output(['hostname', '--fqdn']).rstrip()
apps_for_pkg = {}
f = open('/etc/yandex-direct/direct-apps.conf.yaml', 'r')
conf = yaml.safe_load(f)['apps']
for app in conf:
    pkg = conf[app]['package']
    if pkg not in apps_for_pkg:
        apps_for_pkg[pkg] = []
    apps_for_pkg[pkg].append(app)
f.close()

lf = open('/var/lock/dt-write-versions-to-zk.lock', 'w+')
lockf(lf, LOCK_EX)

dpkg_status_stat = os.stat('/var/lib/dpkg/status')

zk_hosts = ['ppcback01f.yandex.ru:2181', 'ppcback01i.yandex.ru:2181', 'ppcback01e.yandex.ru:2181']
zk = KazooClient(hosts=','.join(zk_hosts))
zk.start()
dpkg_query_output = check_output(['dpkg-query', '--show', '--showformat', '${Status}\t${Package}\t${version}\n'])
for line in dpkg_query_output.rstrip().split('\n'):
    status, pkg, version = line.split('\t')
    if status == 'install ok installed' and pkg in apps_for_pkg:
        for app in apps_for_pkg[pkg]:
            zk_node = zk_node_format.format(app=app, fqdn=fqdn)
            zk.ensure_path(zk_node)
            time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            data = '%s\ntime=%s\n' % (version, time)
            zk.set(zk_node, data)

if mark_file_path:
    if not os.path.exists(mark_file_path):
        open(mark_file_path, 'w+').close()
    os.utime(mark_file_path, (dpkg_status_stat.st_atime, dpkg_status_stat.st_mtime))

logging.info('FINISH')
