import json
import time
import os
import sys
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed

BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(BASE_DIR)

from aidso_geo.models.process import process_success, platform_process, scheduler, task_send_queue
from aidso_geo.utils.bh_utils import query_data
from loguru import logger


def safe_load(v):
    if not v:
        return v
    if isinstance(v, (dict, list)):
        return v
    try:
        return json.loads(v)
    except Exception:
        return v


def preprocess_item(item):
    """
    统一处理需要反序列化的字段
    """
    if item.get("comWordsMap"):
        item["comWordsMap"] = safe_load(item.get("comWordsMap"))
    if item.get("brandWords"):
        item["brandWords"] = safe_load(item.get("brandWords"))
    if item.get("comWords"):
        item["comWords"] = safe_load(item.get("comWords"))
    if item.get("keywords"):
        item["keywords"] = safe_load(item.get("keywords"))
    if item.get("productWordsMap"):
        item["productWordsMap"] = safe_load(item.get("productWordsMap"))
    return item


def log_task(item):
    logger.success(
        f"{item.get('reqId')}---{item.get('platform')}---{item.get('prompt')}---{item.get('type')}"
    )


def run_parallel_task(data_list, handler, max_workers=5):
    """
    并发执行任务
    """
    if not data_list:
        return

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_map = {executor.submit(handler, item): item for item in data_list}

        for future in as_completed(future_map):
            item = future_map[future]
            try:
                future.result()
            except Exception as e:
                logger.exception(
                    f"任务执行失败: taskId={item.get('taskId')}, reqId={item.get('reqId')}, error={e}"
                )


def processing_worker():
    """
    线程1：
    处理 status = PROCESSING 的任务
    """
    while True:
        try:
            process_list = query_data(
                "select * from geo_commit_task "
                "where status = 'PROCESSING' "
                "and insertime > unix_timestamp() - 36000"
            )
            if process_list:
                for idx, item in enumerate(process_list):
                    process_list[idx] = preprocess_item(item)
                    log_task(process_list[idx])
            run_parallel_task(process_list, platform_process, max_workers=10)

        except Exception as e:
            logger.exception(f"processing_worker 主循环异常: {e}")

        time.sleep(180)


def stream_worker():
    """
    线程2：
    处理 type = stream 的任务
    """
    while True:
        try:
            stream_list = query_data(
                "select * from geo_commit_task "
                "where type = 'stream' "
                "and status != 'SUCCESS' "
                "and unix_timestamp() - insertime > 480 "
                "order by insertime desc"
            )

            if stream_list:
                for idx, item in enumerate(stream_list):
                    stream_list[idx] = preprocess_item(item)
                    log_task(stream_list[idx])

                run_parallel_task(stream_list, process_success, max_workers=3)

        except Exception as e:
            logger.exception(f"stream_worker 主循环异常: {e}")

        time.sleep(180)


def stream_batch_worker():
    """
    线程3：
    处理 type = stream_batch 的任务
    """
    while True:
        try:
            stream_batch_list = query_data(
                "select * from geo_commit_task "
                "where type = 'stream_batch' "
                "and status != 'SUCCESS' "
                "and unix_timestamp() - insertime > 18000 "
                "order by insertime desc"
            )

            # stream_batch_list = query_data("select * from geo_commit_task where taskId = '000aa9c0-25f0-4f87-84cb-95309d34bb12' and platform = 'TYQWA'")

            if stream_batch_list:
                for idx, item in enumerate(stream_batch_list):
                    stream_batch_list[idx] = preprocess_item(item)
                    log_task(stream_batch_list[idx])

                run_parallel_task(stream_batch_list, process_success, max_workers=3)

        except Exception as e:
            logger.exception(f"stream_batch_worker 主循环异常: {e}")

        time.sleep(180)


def batch_worker():

    while True:
        try:
            batch_list = query_data(
                "select * from geo_commit_task "
                "where type = 'batch' "
                "and status != 'SUCCESS' "
                "and unix_timestamp() - insertime > 60000 "
                "order by insertime desc"
            )

            if batch_list:
                for idx, item in enumerate(batch_list):
                    batch_list[idx] = preprocess_item(item)
                    log_task(batch_list[idx])

                run_parallel_task(batch_list, process_success, max_workers=10)

        except Exception as e:
            logger.exception(f"batch_worker 主循环异常: {e}")

        time.sleep(180)


def main():
    thread_list = [
        # threading.Thread(target=processing_worker, name="processing_worker", daemon=True),
        threading.Thread(target=stream_worker, name="stream_worker", daemon=True),
        threading.Thread(target=stream_batch_worker, name="stream_batch_worker", daemon=True),
        threading.Thread(target=batch_worker, name="batch_worker", daemon=True),
    ]

    for t in thread_list:
        t.start()
        logger.info(f"线程启动成功: {t.name}")

    while True:
        for t in thread_list:
            if not t.is_alive():
                logger.error(f"线程已退出: {t.name}")
        time.sleep(60)


if __name__ == "__main__":
    main()


