import pandas as pd
from pymongo import MongoClient
import pickle
from flask import Flask,request
import time
import logging
import traceback

app = Flask('model_prediction_lightgbm——service')


def get_data_from_mongo(args):
    mongodb_connection,mongodb_database,mongodb_read_table,timeBegin,timeEnd = "mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/",args['mongodb_database'],args['mongodb_read_table'],args['timeBegin'],args['timeEnd']
    client = MongoClient(mongodb_connection)
    # 选择数据库（如果数据库不存在，MongoDB 会自动创建）
    db = client[mongodb_database]
    collection = db[mongodb_read_table]  # 集合名称
    query = {"dateTime": {"$gte": timeBegin, "$lte": timeEnd}}
    cursor = collection.find(query)
    data = list(cursor)
    df = pd.DataFrame(data)
    # 4. 删除 _id 字段（可选）
    if '_id' in df.columns:
        df = df.drop(columns=['_id'])
    client.close()
    return df
    

def insert_data_into_mongo(res_df,args):
    mongodb_connection,mongodb_database,mongodb_write_table = "mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/",args['mongodb_database'],args['mongodb_write_table']
    client = MongoClient(mongodb_connection)
    db = client[mongodb_database]
    if mongodb_write_table in db.list_collection_names():
        db[mongodb_write_table].drop()
        print(f"Collection '{mongodb_write_table} already exist, deleted successfully!")
    collection = db[mongodb_write_table]  # 集合名称
    # 将 DataFrame 转为字典格式
    data_dict = res_df.to_dict("records")  # 每一行作为一个字典
    # 插入到 MongoDB
    collection.insert_many(data_dict)
    print("data inserted successfully!")
    

def model_prediction(df,args):
    mongodb_connection,mongodb_database,mongodb_model_table,model_name = "mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/",args['mongodb_database'],args['mongodb_model_table'],args['model_name']
    client = MongoClient(mongodb_connection)
    db = client[mongodb_database]
    collection = db[mongodb_model_table]
    model_data = collection.find_one({"model_name": model_name})
    if model_data is not None:
        model_binary = model_data['model']  # 确保这个字段是存储模型的二进制数据
        # 反序列化模型 
        model = pickle.loads(model_binary)
        df['predict'] = model.predict(df[model.feature_name()])
        print("model predict result  successfully!")
    return df


@app.route('/model_prediction_lightgbm', methods=['POST'])
def model_prediction_lightgbm():
    # 获取程序开始时间  
    start_time = time.time()  
    result = {}
    success = 0
    print("Program starts execution!")
    try:
        args = request.values.to_dict()
        print('args',args)
        logger.info(args)
        power_df = get_data_from_mongo(args)
        model = model_prediction(power_df,args)
        insert_data_into_mongo(model,args)
        success = 1
    except Exception as e:
        my_exception = traceback.format_exc()
        my_exception.replace("\n","\t")
        result['msg'] = my_exception
    end_time = time.time() 
   
    result['success'] = success
    result['args'] = args
    result['start_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))
    result['end_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time))
    print("Program execution ends!")
    return result


if __name__=="__main__":  
    print("Program starts execution!")
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    logger = logging.getLogger("model_prediction_lightgbm log")
    from waitress import serve
    serve(app, host="0.0.0.0", port=10090)
    print("server start!")