import os
import sys
import threading
import time
import traceback
import socket
import uuid

import numpy as np
import pandas as pd
import redis
import logging

os.environ["PYARROW_IGNORE_TIMEZONE"] = "1"
sys.path.append(os.path.dirname(sys.path[0]))  # 上级目录
# from utils.templates import Templates
from sqlalchemy import text
from vgg_model import VGGNet
from utils.db_util import DbTypes, DBUtil, get_redis_h14
logging.basicConfig(format='%(asctime)s %(name)s %(levelname)s %(message)s', level=logging.INFO)


class ImgExtractFeatures(object):

    def __init__(self, site_name='us', img_type="amazon_inv", thread_num=10, limit=1000):
        # super(ImgFeatures, self).__init__()

        self.site_name = site_name
        self.img_type = img_type
        self.thread_num = thread_num
        self.limit = limit
        self.engine_doris = DBUtil.get_db_engine(db_type=DbTypes.doris.name, site_name=self.site_name)
        self.client_redis = get_redis_h14()
        self.local_name = f"{self.site_name}_img_features"
        self.vgg_model = VGGNet()
        self.hostname = socket.gethostname()
        self.read_table = f"img_local_path"
        self.save_table = f"img_features"

    def acquire_lock(self, lock_name, timeout=100):
        """
        尝试获取分布式锁, 能正常设置锁的话返回True, 不能设置锁的话返回None
        lock_name: 锁的key, 建议和任务名称保持一致
        """
        lock_value = str(uuid.uuid4())
        lock_acquired = self.client_redis.set(lock_name, lock_value, nx=True, ex=timeout)  # 可以不设置超时时间
        # lock_acquired = self.client_redis.set(lock_name, lock_value, nx=True)
        return lock_acquired, lock_value

    def release_lock(self, lock_name, lock_value):
        """释放分布式锁"""
        script = """
        if redis.call("get", KEYS[1]) == ARGV[1] then
            return redis.call("del", KEYS[1])
        else
            return 0
        end
        """
        result = self.client_redis.eval(script, 1, lock_name, lock_value)
        return result

    def read_data(self):
        while True:
            try:
                lock_acquired, lock_value = self.acquire_lock(lock_name=self.local_name)
                if lock_acquired:
                    print("self.hostname:", self.hostname)
                    with self.engine_doris.begin() as conn:
                        sql_read = text(f"SELECT id, img_unique, local_path, img_type FROM selection.{self.read_table} WHERE site_name='{self.site_name}' and img_type='{self.img_type}' and state=1 LIMIT {self.limit};")
                        # result = conn.execute(sql_read)
                        # df = pd.DataFrame(result.fetchall())
                        df = pd.read_sql(sql=sql_read, con=self.engine_doris)
                        img_unique_tuple = tuple(df.img_unique)
                        print(f"sql_read: {sql_read}, {df.shape}", img_unique_tuple[:10])
                        if img_unique_tuple:
                            img_unique_tuple_str = f"('{img_unique_tuple[0]}')" if len(img_unique_tuple) == 1 else f"{img_unique_tuple}"
                            sql_update = text(f"UPDATE selection.{self.read_table} SET state=2 WHERE img_unique IN {img_unique_tuple_str};")
                            print("sql_update:", sql_update)
                            conn.execute(sql_update)
                    self.release_lock(lock_name=self.local_name, lock_value=lock_value)
                    return df
                else:
                    print(f"当前有其它进程占用redis的锁, 等待5秒继续获取数据")
                    time.sleep(5)  # 等待5s继续访问锁
                    continue
            except Exception as e:
                print(f"读取数据错误: {e}", traceback.format_exc())
                time.sleep(5)
                self.engine_doris = DBUtil.get_db_engine(db_type=DbTypes.doris.name, site_name=self.site_name)

                continue

    def handle_data(self, df, thread_id):
        id_list = list(df.id)
        img_unique_list = list(df.img_unique)
        local_path_list = list(df.local_path)
        data_list = []
        for id, img_unique, local_path in zip(id_list, img_unique_list, local_path_list):
            index = id_list.index(id)
            print(f"thread_id, index, id, img_unique, local_path: {thread_id, index, id, img_unique, local_path}")
            if self.hostname not in ['hadoop5', 'hadoop6', 'hadoop7', 'hadoop8']:
                local_path = local_path.replace("/mnt", "/home")
            try:
                features = self.vgg_model.vgg_extract_feat(file=local_path)
            except Exception as e:
                print(e, traceback.format_exc())
                features = list(np.zeros(shape=(512,)))
            data_list.append([id, img_unique, str(features), self.img_type, self.site_name])
        columns = ['id', 'img_unique', 'features', 'img_type', 'site_name']
        df_save = pd.DataFrame(data_list, columns=columns)
        return df_save

    def save_data(self, df):
        df.to_sql(self.save_table, con=self.engine_doris, if_exists="append", index=False)
        with self.engine_doris.begin() as conn:
            img_unique_tuple = tuple(df.img_unique)
            if img_unique_tuple:
                img_unique_tuple_str = f"('{img_unique_tuple[0]}')" if len(img_unique_tuple) == 1 else f"{img_unique_tuple}"
                sql_update = f"update selection.{self.read_table} set state=3 where img_unique in {img_unique_tuple_str};"
                print(f"sql_update: {sql_update}")
                conn.execute(sql_update)

    def run(self, thread_id=1):
        while True:
            try:
                df = self.read_data()
                if df.shape[0]:
                    df_save = self.handle_data(df=df, thread_id=thread_id)
                    self.save_data(df=df_save)
                    # break
                else:
                    break
            except Exception as e:
                print(e, traceback.format_exc())
                self.engine_doris = DBUtil.get_db_engine(db_type=DbTypes.doris.name, site_name=self.site_name)
                self.client_redis = get_redis_h14()
                self.vgg_model = VGGNet()
                time.sleep(20)
                continue

    def run_thread(self):
        thread_list = []
        for thread_id in range(self.thread_num):
            thread = threading.Thread(target=self.run, args=(thread_id, ))
            thread_list.append(thread)
            thread.start()
        for thread in thread_list:
            thread.join()
        logging.info("所有线程处理完成")


if __name__ == '__main__':
    # handle_obj = PicturesFeatures(self_flag='_self')
    # site_name = int(sys.argv[1])  # 参数1：站点
    # site_name = 'us'
    # img_type = "amazon_inv"
    # limit = 100
    # thread_num = 1
    site_name = sys.argv[1]  # 参数1：站点
    img_type = sys.argv[2]  # 参数2：图片来源类型
    limit = int(sys.argv[3])  # 参数3：每次读取的数量--1000
    thread_num = int(sys.argv[4])  # 参数4：线程数量--5
    handle_obj = ImgExtractFeatures(site_name=site_name, img_type=img_type, thread_num=thread_num, limit=limit)
    # handle_obj.run()
    handle_obj.run_thread()