<reponame>qbzenker/vearch
# Copyright 2019 The Vearch Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# ==============================================================================
import json
import sys
import time
import os
import requests
import subprocess

import config
import exceptions


def installFFmpeg():
    """install ffmpeg method"""
    p = subprocess.call('bash ../bin/centos_yum_install_ffmpeg.sh', shell=True)
    if p == 0:
        raise exceptions.InstallError('FFmpeg install failed!')
    print('FFmpeg install success!')


def create_DB_space():
    db = dict(name=config.video['db'])
    space = dict(name=config.video['space'],
                 dynamic_schema='strict',
                 partition_num=2,
                 replica_num=1,
                 engine=dict(name='gamma',
                             metric_type='InnerProduct'),
                 properties=dict(url=dict(type='keyword',
                                          index='true'),
                                 feature=dict(type='vector',
                                              dimension=512,
                                              format='normalization'
                                              )
                                 )
                 )
    headers = {'content-type':'application/json'}
    ip = '%s/db/_create'%(config.video['ip'])
    res = requests.put(ip, data=json.dumps(db), headers=headers)
    ip = '%s/space/%s/_create' % (config.video['ip'], db['name'])
    res = requests.put(ip, data=json.dumps(db), headers=headers)
    return res.json()


def run():
    main_p = subprocess.Popen(f'python main.py  --model_name=face_retrieval', shell=True)
    process_list.append(main_p)
    # wait 10s for main_p load model
    time.sleep(10)
    res = create_DB_space()
    if res['code'] in [200, 550, 561]:
        print(res)
    else:
        raise exceptions.CreateDBAndSpaceError()

    imgpath = config.video['imagepath']
    isExists = os.path.exists(imgpath)
    path = config.video['videopath']
    if not isExists:
        os.makedirs(imgpath)
    ffmpeg_p = subprocess.Popen(
        f"ffmpeg -i {path} -vf select='eq(pict_type\,I)' -vsync 2 -f image2 {imgpath} img-%d.jpg", shell=True)
    process_list.append(ffmpeg_p)
    count = 1
    insert_ip = f'%s:%d/%s/%s'%(config.video['ip'], config.port, config.video['db'], config.video['space'])
    insert_data = dict(url=None, feature=dict(feature=None))
    headers = {'content-type':'application/json'}
    while True:
        if main_p.poll() is not None:
            raise Exception('main process is killed!')
        if ffmpeg_p.poll() is not None:
            raise Exception('ffmpeg process is killed!')
        imagefile = imgpath + f'/img-{count}.jpg'
        if os.path.exists(imagefile):
            insert_data['url'] = imagefile
            insert_data['feature']['feature'] = imagefile
            res = requests.post(insert_ip, data=json.dumps(insert_data), headers=headers)
            count += 1
            time.sleep(1)


def clear():
    for p in process_list:
        if p.poll() is None:
            p.terminate()
    sys.exit(1)


if __name__ == '__main__':
    # judge if ffmpeg has been installed
    p = subprocess.run('ffmpeg -version', shell=True)
    if p.returncode != 0:
        installFFmpeg()
    process_list = []
    try:
        run()
    except BaseException as err:
        print(err)
        clear()



