import sys
import os

sys.path.append(os.path.dirname(sys.path[0]))
import multiprocessing
import traceback
from utils.db_connect import BaseUtils


def check_realtime_tasks(site) -> bool:
    """快速检查是否有待处理的实时抓取任务"""
    try:
        engine_pg = BaseUtils(site_name=site).pg_connect()
        sql = f"SELECT id FROM realtime_crawl_asin WHERE site='{site}' AND state=0 LIMIT 1"
        df = engine_pg.read_sql(sql)
        return not df.empty
    except Exception as e:
        print(f'check_realtime_tasks 查询失败: {e}\n{traceback.format_exc()}')
        return False


def _read_realtime_tasks(site, read_size=1500) -> list:
    """读取待抓取任务，原子性标记为 state=2，返回 id|asin|task_id 列表"""
    try:
        engine_pg = BaseUtils(site_name=site).pg_connect()
        sql_read = (
            f"SELECT id, asin, task_id FROM realtime_crawl_asin "
            f"WHERE site='{site}' AND state=0 "
            f"FOR UPDATE SKIP LOCKED LIMIT {read_size}"
        )
        df = engine_pg.read_then_update(
            select_sql=sql_read,
            update_table='realtime_crawl_asin',
            set_values={'state': 2},
            where_keys=['id'],
        )
        if df.empty:
            return []
        return list(df['id'].astype(str) + '|' + df['asin'] + '|' + df['task_id'].astype(str))
    except Exception as e:
        print(f'_read_realtime_tasks 读取失败: {e}\n{traceback.format_exc()}')
        return []


def _realtime_task_worker(site, data_chunk):
    """子进程入口：实例化爬虫并执行"""
    print(f"实时任务子进程 {os.getpid()} 分配到 {len(data_chunk)} 条ASIN")
    from amazon_spider.realtime_report_spider import RealtimeReportSpider
    task_id = data_chunk[0].split('|')[2] if data_chunk else None
    spider = RealtimeReportSpider(site_name=site, task_id=task_id)
    spider.run_with_data(data_chunk)


def run_realtime_spider(site, num_processes=5):
    """读取实时任务并启动子进程抓取"""
    all_data = _read_realtime_tasks(site)
    if not all_data:
        print(f'run_realtime_spider: {site} 无待处理任务')
        return
    print(f'实时任务 {site}: 共 {len(all_data)} 条ASIN，启动子进程处理')
    # 与月表爬虫保持一致：数据少于300条时动态缩减进程，否则全量5进程
    if len(all_data) < 300:
        actual_processes = max(1, len(all_data) // 60)
    else:
        actual_processes = num_processes
    chunks = [all_data[i::actual_processes] for i in range(actual_processes)]
    processes = []
    for chunk in chunks:
        if chunk:
            p = multiprocessing.Process(target=_realtime_task_worker, args=(site, chunk))
            processes.append(p)
            p.start()
    for p in processes:
        p.join()
    print(f'实时任务 {site}: 全部子进程完成')
