print('抓取 月 asin')
import os
import sys
from datetime import datetime, timedelta

sys.path.append(os.path.dirname(sys.path[0]))
import time
import random
import multiprocessing
from amazon_spider.asin_detail_pg import async_asin_pg
from amazon_spider.VPS_IP import pppoe_ip
from amazon_save_db.save_asin_detail_pg import Save_asin_detail
from threading_spider.Poll_site_spider import start_day_spider
from threading_spider.db_connectivity import connect_db
from threading_spider.post_to_dolphin import DolphinschedulerHelper
from utils.secure_db_client import get_remote_engine
from utils.db_connect import BaseUtils


def db_engine(site_name, db_type):
    return get_remote_engine(site_name=site_name, db_type=db_type)


def select_asin_state(site, date_info):
    """检查 PG {site}_all_syn_st_month_{年}_{月} 是否还有 state in (1,2)。有→[1]，无→None"""
    table = f"{site}_all_syn_st_month_{date_info.replace('-', '_')}"
    for _ in range(5):
        try:
            sql = f"SELECT id FROM {table} where state in (1,2) LIMIT 1"
            engine_pg = db_engine(site, 'postgresql_14_outer')
            df = engine_pg.read_sql(sql)
            return [1] if not df.empty else None
        except Exception as e:
            print(e, '查询PG state报错', sql)
            time.sleep(3)
    return [1]


def db_cursor_connect_update(sql, site=None):
    for _ in range(3):
        try:
            engine = db_engine('us', 'mysql')
            print('更新sql：', sql)
            with engine.begin() as conn:
                conn.execute(sql)
            return
        except Exception as e:
            print(site, 'db_cursor_connect 报错：', e, sql)


def has_pending_asin_and_mark_started(site, date_info):
    """
    入口：
    1) PG 没ASIN → False，进程退出
    2) PG 有ASIN → workflow_manager state=1→2 推进（行锁天然幂等）
    """
    if select_asin_state(site, date_info) is None:
        print(f'[启动] {site} {date_info} PG 里没ASIN要抓，退出')
        return False

    sql_update = (f"UPDATE workflow_manager set spider_state=2 "
                  f"WHERE site_name='{site}' and date_type='month' and date_info='{date_info}' "
                  f"and spider_is_ready='yes' and spider_name='{site}_spider_asin' and spider_state=1")
    db_cursor_connect_update(sql_update, site)
    print(f'[启动] {site} {date_info} 进入抓取循环')
    return True


def trigger_kafka_realtime_dolphin(site, date_info):
    """状态检查 + NX 锁双保险：kafka_state 必须 = 1 才允许触发，确保 dolphin 只触发 1 次"""
    redis_client = BaseUtils().redis_db()
    db_class = connect_db(site)

    # 第 1 道：状态检查（kafka_state 必须是 1 才允许触发）
    try:
        sql_check = (f"SELECT kafka_state FROM workflow_manager "
                     f"WHERE site_name='{site}' and date_type='month' and date_info='{date_info}' "
                     f"and spider_name='{site}_spider_asin'")
        df = db_engine('us', 'mysql').read_sql(sql_check)
        if df.empty:
            print(f'[kafka] {site} {date_info} workflow_manager 没有匹配行，跳过')
            return False
        current_state = df.kafka_state.iloc[0]
        if current_state != 1:
            print(f'[kafka] {site} {date_info} kafka_state={current_state} 已触发过，跳过')
            return False
    except Exception as e:
        print('[kafka] 查 kafka_state 异常', e)
        return False

    # 第 2 道：NX 锁（80 台同时启动时防瞬时撞车）
    nx_key = f"asin_kafka_realtime_{site}_{date_info}"
    if not redis_client.set(nx_key, 1, nx=True, ex=60):
        print(f'[kafka] {site} {date_info} NX 已锁，跳过')
        return False

    # 推进 kafka_state=1→2 + 触发 dolphin + 发消息
    sql = (f"UPDATE workflow_manager set kafka_state=2 "
           f"WHERE site_name='{site}' and date_type='month' and date_info='{date_info}' "
           f"and spider_name='{site}_spider_asin' and kafka_state=1")
    db_cursor_connect_update(sql, site)

    DolphinschedulerHelper.start_process_instance_common(
        project_name="big_data_selection",
        process_df_name='ALL-实时-插件+流量选品',
        startParams={"site_name": site, "date_type": "month", "date_info": date_info},
        warning_Type="ALL"
    )
    db_class.send_mg(
        'pengyanbing,chenyuanjie,fangxingjun,zhouyuchen,guohuanjie',
        f'{site} 开启 kafka 实时消费 {date_info}',
        f'{site} 开始调度 ALL-实时-插件+流量选品'
    )
    print(f'[kafka] {site} {date_info} 已触发 ALL-实时-插件+流量选品')
    return True


def wait_kafka_ready(site, date_info, max_wait_min=40):
    """所有机器都要等 kafka_state=3。超时返回 False。"""
    interval_sec = 120
    max_loops = max_wait_min * 60 // interval_sec   # 40 分钟 = 20 次
    sql = (f"SELECT id FROM workflow_manager "
           f"WHERE site_name='{site}' and date_type='month' and date_info='{date_info}' "
           f"and spider_name='{site}_spider_asin' and kafka_state=3")
    for n in range(max_loops):
        try:
            df = db_engine('us', 'mysql').read_sql(sql)
            if not df.empty:
                print(f'[kafka] {site} {date_info} 已就绪 ({n+1}次查询后)')
                return True
        except Exception as e:
            print('查 kafka_state 报错', e)
        print(f'[kafka] {site} 等待中 ({n+1}/{max_loops})')
        time.sleep(interval_sec)

    print(f'[kafka] {site} {date_info} 等待 {max_wait_min} 分钟超时')
    return False


def finalize_and_trigger_dolphin(site, date_info):
    """结尾：state=2→3 + Redis NX 锁 + 触发 ALL站点-启动30day/月流程"""
    db_class = connect_db(site)
    redis_client = BaseUtils().redis_db()
    month = date_info.split('-')[1]

    # 第 1 道：PG 还有未抓数据 → 不触发
    if select_asin_state(site, date_info) is not None:
        print('PG 还有未完成的，暂不触发调度')
        return False

    # 第 2 道：workflow_manager spider_state 必须 = 2 才允许触发
    try:
        sql_check = (f"SELECT spider_state FROM workflow_manager "
                     f"WHERE site_name='{site}' and date_type='month' and date_info='{date_info}' "
                     f"and spider_name='{site}_spider_asin'")
        df = db_engine('us', 'mysql').read_sql(sql_check)
        if df.empty:
            print(f'[调度] {site} {date_info} workflow_manager 没匹配行，跳过')
            return False
        current_state = df.spider_state.iloc[0]
        if current_state != 2:
            print(f'[调度] {site} {date_info} spider_state={current_state} 已触发过，跳过')
            return False
    except Exception as e:
        print('[调度] 查 spider_state 异常', e)
        return False

    # 第 3 道：NX 锁
    nx_key = f"asin_dolphin_scheduled_{site}_{date_info}"
    if not redis_client.set(nx_key, 1, nx=True, ex=1800):
        print(f'[调度] {site} {date_info} NX 已锁，跳过')
        return False

    ii = 0
    id_tuple = None
    for _ in range(5):
        id_tuple = select_asin_state(site, date_info)
        time.sleep(180)
        if id_tuple is None:
            ii += 1
            if ii > 3:
                break

    if id_tuple is not None:
        redis_client.delete(nx_key)
        print(f'[调度] {site} 复核又有数据，归还NX，本次不触发')
        return False

    sql_done = (f"UPDATE workflow_manager set spider_state=3, spider_is_ready='no' "
                f"WHERE site_name='{site}' and date_type='month' and date_info='{date_info}' "
                f"and spider_is_ready='yes' and spider_name='{site}_spider_asin' and spider_state=2")
    db_cursor_connect_update(sql_done, site)

    DolphinschedulerHelper.start_process_instance_common(
        project_name="big_data_selection",
        process_df_name='ALL站点-启动30day/月流程',
        startParams={"site_name": site, "date_type": "month", "date_info": date_info}
    )
    account = 'pengyanbing,chenyuanjie,fangxingjun,zhouyuchen,guohuanjie'
    title = site + '站点 asin 详情'
    content = f'{month} 月 asin 详情已结束，成功调度 ALL站点-启动30day/月流程'
    db_class.send_mg(account, title, content)
    print(f'[调度] {site} {date_info} 已成功触发')
    return True


def read_data_from_pg(save_detail_reader):
    """主进程：复用 Save_asin_detail.read_db_data()，保留 minid_maxid 分段读取"""
    asin_6field = save_detail_reader.read_db_data()
    if not asin_6field:
        return []
    df = save_detail_reader.df_read
    all_data = list(
        df.id.astype("U") + '|' + df.asin + '|' + df.date_info + '|' +
        df.asin_is_variation.astype("U") + '|' + df.data_type.astype("U") + '|' +
        df.volume.astype("U") + '|' + df.weight_str.astype("U"))
    print(f'主进程读取 {len(all_data)} 条ASIN')
    return all_data


def long_time_task(site, date_info, spider_int, data_chunk):
    """子进程：接收分发数据，不自己读数据库"""
    print(f"当前进程 {os.getpid()} 分配到 {len(data_chunk)} 条数据")
    print(f"开始=============== {site} asin 抓取 =======================")
    async_asin_us = async_asin_pg(site_name=site, date_info=date_info, spider_int=spider_int, skip_read_init=True)
    async_asin_us.run_with_data(data_chunk)


if __name__ == '__main__':
    # 执行命令 /usr/local/bin/python3 /mnt/py_spider/threading_spider/Poll_site_spider_month.py us month 2026-04
    pppoe_ip()
    site = sys.argv[1]                    # 'us'
    date_type = sys.argv[2]               # 'month'
    date_info = sys.argv[3]               # '2026-04'
    month = date_info.split('-')[1]       # '04'
    spider_int = 0
    num_processes = 5
    print(f'站点: {site}, date_type: {date_type}, date_info: {date_info}, month: {month}')

    if not has_pending_asin_and_mark_started(site, date_info):
        sys.exit(0)

    # 单台触发 kafka 实时消费 dolphin（其余 49 台 NX 失败但继续等）
    trigger_kafka_realtime_dolphin(site, date_info)

    # 所有机器等 kafka 就绪，超时退出本站点（不抓取）
    if not wait_kafka_ready(site, date_info):
        print(f'[启动] {site} {date_info} kafka 等待超时，退出本站点')
        sys.exit(0)

    reader = Save_asin_detail(site_name=site, date_info=date_info, spider_int=spider_int)
    reader.read_size = 1500
    empty_count = 0
    day_spider_ran_date = None

    while True:
        # 凌晨2-5点切换日ASIN任务
        _today = datetime.now().strftime('%Y-%m-%d')
        if 2 <= datetime.now().hour <= 5 and day_spider_ran_date != _today:
            day_spider_ran_date = _today
            print(f'==== {datetime.now()} 凌晨2点，执行日ASIN抓取 ====')
            try:
                start_day_spider()
            except Exception as e:
                print(f'日ASIN任务异常: {e}')
            print(f'==== {datetime.now()} 日ASIN任务完成，继续月度任务 ====')

        five_minutes_later = datetime.now() + timedelta(minutes=2)
        print('后五分钟时间', five_minutes_later)

        all_data = read_data_from_pg(reader)
        if not all_data:
            empty_count += 1
            if reader.minid_maxid_list:
                print('分段区间无数据，清除分段限制重试')
                reader.minid_maxid_list = []
                continue
            if empty_count >= 3:
                print(f'连续{empty_count}次读空，检查PG状态')
                if select_asin_state(site, date_info) is None:
                    if finalize_and_trigger_dolphin(site, date_info):
                        break  # winning 完成
                    # finalize 返回 False：要么 NX 被别人拿（49台），要么复核期数据回流（winning）
                    if select_asin_state(site, date_info) is not None:
                        # winning 检测到回流 → 单机继续抓
                        print('数据回流，单机继续抓取')
                        empty_count = 0
                        continue
                    # 49 台 NX 失败 + PG 还是空 → break 退出
                    print('其他机器已在处理 finalize，本机退出')
                    break
                time.sleep(60)
                empty_count = 0
                continue
            time.sleep(10)
            continue

        empty_count = 0
        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 idx in range(actual_processes):
            if chunks[idx]:
                p = multiprocessing.Process(target=long_time_task, args=(site, date_info, spider_int, chunks[idx]))
                processes.append(p)
                p.start()
        for p in processes:
            p.join()

        print('所有进程运行完毕！开始切换ip')
        if datetime.now() > five_minutes_later:
            pppoe_ip()
        else:
            time.sleep(random.uniform(150, 220))
            pppoe_ip()
