import json
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

from loguru import logger
from apscheduler.schedulers.blocking import BlockingScheduler

from aidso_geo.config.base_config import init_redis
from aidso_geo.models.process import main_process


redis_client = init_redis()

QUEUE_KEY = "geo:task_commit:list"

# 每轮最多拉取任务数量
BATCH_SIZE = 1000

# 每 60 秒执行一次
INTERVAL_SECONDS = 60

# 每轮内部并发数
CONCURRENT_WORKERS = 40


def parse_task(raw):
    try:
        if isinstance(raw, bytes):
            raw = raw.decode("utf-8")

        return json.loads(raw)

    except Exception as e:
        logger.exception(f"解析队列数据失败: {e}, raw={raw}")
        return None


def pull_tasks():
    """
    每次最多从 Redis 队列拉取 BATCH_SIZE 条任务
    """
    task_list = []

    for _ in range(BATCH_SIZE):
        raw = redis_client.rpop(QUEUE_KEY)

        if not raw:
            break

        data = parse_task(raw)

        if data:
            task_list.append(data)

    return task_list


def handle_one_task(task):
    """
    单条任务处理。
    成功：正常结束
    失败：重新放回 Redis 队列
    """
    req_id = task.get("reqId", "")
    platform = task.get("platform", "")
    prompt = task.get("prompt", "")

    try:
        main_process(task)
        return True

    except Exception as e:
        logger.exception(f"{req_id}--{platform}--处理失败，重新放回队列: {e}")

        redis_client.lpush(
            QUEUE_KEY,
            json.dumps(task, ensure_ascii=False)
        )

        return False


def consume_task_queue():
    """
    定时任务：
    1. 拉取 Redis 队列
    2. 使用线程池并发执行 main_process
    """
    start_time = time.time()

    try:
        tasks = pull_tasks()

        if not tasks:
            logger.info("本轮没有任务")
            return

        with ThreadPoolExecutor(max_workers=CONCURRENT_WORKERS) as executor:
            futures = [
                executor.submit(handle_one_task, task)
                for task in tasks
            ]

            for future in as_completed(futures):
                try:
                    ok = future.result()

                except Exception as e:
                    logger.exception(f"线程任务异常: {e}")

        cost = round(time.time() - start_time, 2)

        logger.success(
            f"本轮处理完成，总数: {len(tasks)}, "
            f"耗时: {cost}s"
        )

    except Exception as e:
        logger.exception(f"定时消费任务异常: {e}")


if __name__ == "__main__":
    scheduler = BlockingScheduler(timezone="Asia/Shanghai")

    scheduler.add_job(
        consume_task_queue,
        trigger="interval",
        seconds=INTERVAL_SECONDS,
        id="consume_geo_task_commit_queue",
        max_instances=5,       # 允许最多 5 个调度批次同时跑
        coalesce=False,        # 不合并错过的执行
        misfire_grace_time=60
    )

    logger.success(
        f"geo task commit consumer 启动，每 {INTERVAL_SECONDS}s 执行一次，"
        f"QUEUE_KEY={QUEUE_KEY}, "
        f"BATCH_SIZE={BATCH_SIZE}, "
        f"CONCURRENT_WORKERS={CONCURRENT_WORKERS}, "
        f"max_instances=5"
    )

    consume_task_queue()

    scheduler.start()
