from gevent import monkey; monkey.patch_all()
import os
from flask import Flask, request, jsonify
from flask_apscheduler import APScheduler
from gevent.pywsgi import WSGIServer

import config
from hnsw_model import HNSW
from hnsw_model_train import HNSW_Model_Train
from restart_server import restart_retrieval_app
from log_config import LogConfig

app = Flask(__name__)

class APS_Config(object):
    SCHEDULER_API_ENABLED = True
scheduler = APScheduler()

# 定时重启retrieval_app服务, 防止内存累加
@scheduler.task('cron', id='restart_retrieval_app', day='*', hour='02', minute='00', second='00', timezone='Asia/Shanghai')
def retrieval_app_schedule():
    # 重启retrieval_app服务
    restart_retrieval_app()

# 定时训练HNSW模型并重启服务(0-'mon',...2-'wed',...,6-'sun')
@scheduler.task('cron', id='hm_train', week='*', day_of_week='0', hour='02', minute='30', second='00', timezone='Asia/Shanghai')
def hm_train_schedule():
    hm_train = HNSW_Model_Train(hm_logger)
    hm_train()
    os.popen("nohup python restart_server.py 1 > logs/temp_app.log 2>&1 &")

# 每周定时刷新日志
@scheduler.task('cron', id='log_reset', week='*', day_of_week='sun', hour='05', minute='00', second='00', timezone='Asia/Shanghai')
def log_reset_schedule():
    hm_LogConfig.log_reset()
    os.popen("nohup python restart_server.py > logs/temp_app.log 2>&1 &")

# hnsw模型数据检索
@app.route('/retrieve', methods=['GET', 'POST'])
def retrieve():
    if request.method == 'POST':
        # 获取post数据
        query_vec = request.get_json()
        # HNSW检索
        query_labels = hnsw.retrieve(query_vec)
        return jsonify(query_labels)


if __name__ == '__main__':
    # 日志采集初始化
    hm_LogConfig = LogConfig(config.retrieval_path, "hnsw_model")
    hm_logger = hm_LogConfig.get_log()
    # HNSW模型初始化
    hnsw = HNSW(hm_logger)

    # 定时更新HNSW模型增/改/删变化
    app.config.from_object(APS_Config())
    scheduler.init_app(app)
    scheduler.start()
    # app.run(host='0.0.0.0',port='8836')
    server = WSGIServer(('0.0.0.0', 8836), app)
    server.serve_forever()