import os
import sys
from datetime import datetime, timedelta

sys.path.append(os.path.dirname(sys.path[0]))
import multiprocessing
import time
import random
from amazon_spider.day_asin_spider import async_asin_pg
from amazon_spider.VPS_IP import pppoe_ip
from utils.db_connect import BaseUtils
from utils.secure_db_client import get_remote_engine
from threading_spider.post_to_dolphin import DolphinschedulerHelper


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


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 select_asin_state(site, next_date):
    """检查 PG {site}_all_syn_st_day_{年_月_日} 是否还有 state in (1,2)。有→[1]，无→None"""
    table = f"{site}_all_syn_st_day_{next_date.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 has_pending_asin_and_mark_started(site, next_date):
    """
    入口：
    1) workflow_manager 没有 spider_is_ready='yes' 的记录 → False
    2) PG 没ASIN → False
    3) 有 → spider_state=1→2 推进
    """
    sql_check = (f"SELECT id FROM workflow_manager "
                 f"WHERE site_name='{site}' and date_type='day' and date_info='{next_date}' "
                 f"and spider_name='{site}_spider_day_asin' and spider_is_ready='yes'")
    try:
        df = db_engine('us', 'mysql').read_sql(sql_check)
        if df.empty:
            print(f'[启动] {site} {next_date} workflow_manager 无ready记录，跳过')
            return False
    except Exception as e:
        print('[启动] 查 workflow_manager 异常', e)
        return False

    if select_asin_state(site, next_date) is None:
        print(f'[启动] {site} {next_date} PG 里没ASIN要抓，退出')
        return False

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


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

    try:
        sql_check = (f"SELECT kafka_state FROM workflow_manager "
                     f"WHERE site_name='{site}' and date_type='day' and date_info='{next_date}' "
                     f"and spider_name='{site}_spider_day_asin'")
        df = db_engine('us', 'mysql').read_sql(sql_check)
        if df.empty:
            print(f'[kafka] {site} {next_date} workflow_manager 没有匹配行，跳过')
            return False
        current_state = df.kafka_state.iloc[0]
        if current_state != 1:
            print(f'[kafka] {site} {next_date} kafka_state={current_state} 已触发过，跳过')
            return False
    except Exception as e:
        print('[kafka] 查 kafka_state 异常', e)
        return False

    nx_key = f"asin_kafka_realtime_day_{site}_{next_date}"
    if not redis_client.set(nx_key, 1, nx=True, ex=60):
        print(f'[kafka] {site} {next_date} NX 已锁，跳过')
        return False

    sql = (f"UPDATE workflow_manager set kafka_state=2 "
           f"WHERE site_name='{site}' and date_type='day' and date_info='{next_date}' "
           f"and spider_name='{site}_spider_day_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-实时day-插件+流量选品',
        startParams={"site_name": site, "date_type": "day", "date_info": next_date}
    )
    print(f'[kafka] {site} {next_date} 已触发 ALL-实时day-插件+流量选品')
    return True


def wait_kafka_ready(site, next_date, max_wait_min=40):
    """所有机器都要等 kafka_state=3。超时返回 False。"""
    interval_sec = 120
    max_loops = max_wait_min * 60 // interval_sec
    sql = (f"SELECT id FROM workflow_manager "
           f"WHERE site_name='{site}' and date_type='day' and date_info='{next_date}' "
           f"and spider_name='{site}_spider_day_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} {next_date} 已就绪 ({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} {next_date} 等待 {max_wait_min} 分钟超时')
    return False


def finalize_spider_state(site, next_date):
    """结尾：state=2→3, spider_is_ready='no'。日流程不触发额外 dolphin。"""
    redis_client = BaseUtils().redis_db()

    if select_asin_state(site, next_date) is not None:
        print('PG 还有未完成的，暂不结束')
        return False

    try:
        sql_check = (f"SELECT spider_state FROM workflow_manager "
                     f"WHERE site_name='{site}' and date_type='day' and date_info='{next_date}' "
                     f"and spider_name='{site}_spider_day_asin'")
        df = db_engine('us', 'mysql').read_sql(sql_check)
        if df.empty:
            print(f'[结束] {site} {next_date} workflow_manager 没匹配行，跳过')
            return False
        current_state = df.spider_state.iloc[0]
        if current_state != 2:
            print(f'[结束] {site} {next_date} spider_state={current_state} 已结束过，跳过')
            return False
    except Exception as e:
        print('[结束] 查 spider_state 异常', e)
        return False

    nx_key = f"asin_day_finalize_{site}_{next_date}"
    if not redis_client.set(nx_key, 1, nx=True, ex=1800):
        print(f'[结束] {site} {next_date} NX 已锁，跳过')
        return False

    ii = 0
    id_tuple = None
    for _ in range(5):
        id_tuple = select_asin_state(site, next_date)
        time.sleep(80)
        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='day' and date_info='{next_date}' "
                f"and spider_is_ready='yes' and spider_name='{site}_spider_day_asin' and spider_state=2")
    db_cursor_connect_update(sql_done, site)
    print(f'[结束] {site} {next_date} 已成功结束')
    return True


def read_data_from_pg(site, next_date, read_size=1500):
    """主进程：走代理PG read_then_update（含FOR UPDATE）读取数据"""
    try:
        date_info = next_date.replace('-', '_')
        engine_pg = db_engine(site, 'postgresql_14_outer')
        table = f'{site}_all_syn_st_day_{date_info}'
        sql_read = (f"SELECT asin, id, date_info, asin_is_variation, data_type, volume, weight_str "
                    f"FROM {table} WHERE state=1 FOR UPDATE SKIP LOCKED LIMIT {read_size}")
        print(sql_read)
        df_read = engine_pg.read_then_update(
            select_sql=sql_read,
            update_table=table,
            set_values={"state": 2},
            where_keys=["id"],
        )
        df_read.drop_duplicates(['asin'], inplace=True)
        if df_read.shape[0] > 0:
            df_read['volume'] = df_read['volume'].fillna('null')
            df_read['weight_str'] = df_read['weight_str'].fillna('null')
            asin_list = list(
                df_read.id.astype("U") + '|' + df_read.asin + '|' + df_read.date_info + '|' +
                df_read.asin_is_variation.astype("U") + '|' + df_read.data_type.astype("U") + '|' +
                df_read.volume.astype("U") + '|' + df_read.weight_str.astype("U"))
            print(f'读取 {len(asin_list)} 条ASIN')
            return asin_list
        else:
            print('没有 state=1 的数据了')
            return []
    except Exception as e:
        print(f'读取数据失败: {e}')
        return []


def long_time_task(site, next_date, data_chunk):
    print(f"当前进程 {os.getpid()} 分配到 {len(data_chunk)} 条数据")
    async_asin_us = async_asin_pg(site_name=site, next_date=next_date)
    async_asin_us.run_with_data(data_chunk)


def run_one_date(site, next_date, num_processes=5):
    """对单个日期跑完整四段式流程"""
    if not has_pending_asin_and_mark_started(site, next_date):
        return

    trigger_kafka_realtime_dolphin(site, next_date)

    if not wait_kafka_ready(site, next_date):
        print(f'[启动] {site} {next_date} kafka 等待超时，退出本日期')
        return

    empty_count = 0
    while True:
        five_minutes_later = datetime.now() + timedelta(minutes=2)
        print('后五分钟时间', five_minutes_later)
        date_info = next_date.replace('-', '_')

        all_data = read_data_from_pg(site, next_date, read_size=1500)
        if not all_data:
            empty_count += 1
            if empty_count >= 3:
                print(f'连续{empty_count}次读空，检查PG状态')
                if select_asin_state(site, next_date) is None:
                    if finalize_spider_state(site, next_date):
                        break
                    if select_asin_state(site, next_date) is not None:
                        print('数据回流，单机继续抓取')
                        empty_count = 0
                        continue
                    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, 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()


def start_day_spider():
    pppoe_ip()
    site_list = ['us']
    today = datetime.now()
    date_list = [(today - timedelta(days=i)).strftime('%Y-%m-%d') for i in range(5, -1, -1)]
    print(f'本次日ASIN任务日期范围: {date_list}')
    for site in site_list:
        for next_date in date_list:
            run_one_date(site, next_date)
    return None


if __name__ == '__main__':
    print('抓取 日 asin')
    start_day_spider()
