import os
import sys
from datetime import datetime, timedelta

sys.path.append(os.path.dirname(sys.path[0]))  # 上级目录
import multiprocessing
from amazon_spider.search_term_pg import search_temp_pg
from threading_spider.db_connectivity import connect_db
import time
import random
from amazon_spider.VPS_IP import pppoe_ip
from threading_spider.post_to_dolphin import DolphinschedulerHelper
from utils.secure_db_client import get_remote_engine
from utils.db_connect import BaseUtils
from amazon_params.params import DB_SEARCH_TERM_PARAMS_SPIDER


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


def select_search_term_state(date_info, site):
    """检查 PG 里是否还有 state in (1,2) 的搜索词。有→返回[1]，无→返回None"""
    for i in range(5):
        try:
            sql_read = f"SELECT id FROM {site}_search_term_month_syn where state in (1,2) and date_info='{date_info}' LIMIT 1"
            engine_pg = db_engine(site, 'postgresql_14_outer')
            df = engine_pg.read_sql(sql_read)
            return [1] if not df.empty else None
        except Exception as e:
            print(e, '报错11。', sql_read)
            time.sleep(3)
    return 1


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


def has_pending_search_terms_and_mark_started(site, date_info):
    """
    入口判断：
    1) PG 没搜索词 → 返回 False，进程退出
    2) PG 有搜索词 → 用 read_then_update 原子推进 workflow_manager state=1→2
       拿到行的机器是赢家，发"开启"消息；其它机器跳过消息但继续抓取
    """
    id_tuple = select_search_term_state(date_info, site)
    if id_tuple is None:
        print(f'[启动] {site} {date_info} PG 里没搜索词要抓，退出')
        return False

    engine_mysql = db_engine('us', 'mysql')
    sql_read = (f"SELECT id, spider_state FROM workflow_manager "
                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_st' and spider_state=1 "
                f"FOR UPDATE")

    try:
        df = engine_mysql.read_then_update(
            select_sql=sql_read,
            update_table='workflow_manager',
            set_values={"spider_state": 2},
            where_keys=["id"],
        )
    except Exception as e:
        print(f'[启动] 推进 state 异常，跳过消息但继续抓: {e}')
        df = None

    if df is not None and not df.empty and int(df.spider_state.iloc[0]) == 1:
        db_class = connect_db(site)
        month_str = date_info.split('-')[1]
        account = 'pengyanbing,fangxingjun,guohuanjie,chenyuanjie'
        title = site + '站点 搜索词'
        content = f'开启 抓取 {site} {month_str} 月 搜索词'
        db_class.send_mg(account, title, content)
        print(f'[启动] {site} {date_info} 推进 state=1→2，已发送开启消息')
    else:
        print(f'[启动] {site} {date_info} 跳过开启消息 (df 空 / 异常 / state 非 1)')

    print(f'[启动] {site} {date_info} 进入抓取循环')
    return True


def finalize_and_trigger_dolphin(site, month, date_info):
    """
    结尾：state=2→3 + 触发调度。靠 Redis SETNX 保证只有一台机器执行。
    复核失败时归还 NX，让真正抓完的机器有机会再触发。
    """
    db_class = connect_db(site)
    redis_client = BaseUtils().redis_db()

    # 第 1 道：PG 还有未抓数据 → 不触发
    if select_search_term_state(date_info, site) is not None:
        print('搜索词还有未完成的，暂不触发调度')
        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_st'")
        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"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 i in range(5):
        id_tuple = select_search_term_state(date_info, site)
        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_st' and spider_state in (1,2)")
    db_cursor_connect_update(sql_done, site)

    DolphinschedulerHelper.start_process_instance_common(
        project_name="big_data_selection",
        process_df_name='ALL站点-asin同步-pg-api',
        startParams={
            "site_name": f"{site}",
            "date_type": "month",
            "date_info": f'{date_info}'
        }
    )
    account = 'pengyanbing,chenyuanjie,fangxingjun,zhouyuchen,guohuanjie'
    title = site + '站点 搜索词'
    content = f'{month} 月 搜索词 已结束,成功调度 ALL站点-asin同步-pg-api'
    db_class.send_mg(account, title, content)
    print(f'[调度] {site} {date_info} 已成功触发')
    return True


def read_data_from_pg(site, date_info, read_size=1500):
    """主进程：走代理PG read_then_update（含FOR UPDATE）读取数据"""
    try:
        engine_pg = db_engine(site, 'postgresql_14_outer')
        db_search_term = site + DB_SEARCH_TERM_PARAMS_SPIDER["db_search_term"][2:] + '_month_syn'
        sql_read = f"SELECT id, search_term, url FROM {db_search_term} where state=1 and date_info='{date_info}' LIMIT {read_size} for update"
        print(sql_read)
        df_read = engine_pg.read_then_update(
            select_sql=sql_read,
            update_table=db_search_term,
            set_values={"state": 2},
            where_keys=["id"],
        )
        if df_read.shape[0] > 0:
            search_term_list = list(
                df_read.id.astype("U") + '|-|' + df_read.search_term + '|-|' + df_read.url)
            print(f'读取 {len(search_term_list)} 条，已更新state=2')
            return search_term_list
        else:
            print('没有 state=1 的数据了')
            return []
    except Exception as e:
        print(f'读取数据失败: {e}')
        return []


def long_time_task(site, proxy_name, date_info, data_chunk):
    """子进程：用分配到的数据抓取"""
    print(f"当前进程 {os.getpid()} 分配到 {len(data_chunk)} 条数据")
    spider_us = search_temp_pg(site_name=site, read_size=100, proxy_name=proxy_name, date_info=date_info)
    spider_us.run_pol_with_data(data_chunk)


if __name__ == '__main__':
    pppoe_ip()
    site = sys.argv[1]      # 'us'
    date_info = sys.argv[2]  # '2026-04'
    month = date_info.split('-')[1]  # '04'
    num_processes = 5
    proxy_name = None

    print(f'站点: {site}, 月份: {month}, date_info: {date_info}')

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

    while True:
        current_time = datetime.now()
        one_minute_later = current_time + timedelta(minutes=1)
        print('后一分钟时间', one_minute_later)

        all_data = read_data_from_pg(site, date_info, read_size=1800)
        if not all_data:
            print('没有数据，检查是否全部完成')
            if select_search_term_state(date_info, site) is not None:
                print(f'{site} 其他机器还在处理 state=2 数据，等待30s重试')
                time.sleep(30)
                continue
            if finalize_and_trigger_dolphin(site, month, date_info):
                break  # winning 完成
            # finalize 返回 False：要么 NX 被别人拿（49台），要么复核期数据回流（winning）
            if select_search_term_state(date_info, site) is not None:
                # winning 检测到回流 → 单机继续抓
                print('数据回流，单机继续抓取')
                continue
            # 49 台 NX 失败 + PG 还是空 → break 退出
            print('其他机器已在处理 finalize，本机退出')
            break

        chunks = [all_data[i::num_processes] for i in range(num_processes)]

        processes = []
        for idx in range(num_processes):
            if chunks[idx]:
                process = multiprocessing.Process(target=long_time_task, args=(site, proxy_name, date_info, chunks[idx]))
                processes.append(process)
                process.start()
        for process in processes:
            process.join()

        print('所有进程运行完毕！')
        current_time = datetime.now()
        if current_time > one_minute_later:
            pppoe_ip()
        else:
            time.sleep(random.uniform(150, 220))
            pppoe_ip()
