#!/usr/bin/env python3
"""
分布式分红数据拉取 Worker

用法：
    python3 dividend_worker.py <worker_name>

例如：
    python3 dividend_worker.py worker_41

会读取 <worker_name>_tasks.csv，拉取每只股票的分红数据，
输出到 <worker_name>_dividend.csv，并写日志到 <worker_name>.log。
"""

import os
import re
import sys
import time
import warnings
from datetime import datetime

import akshare as ak
import pandas as pd

warnings.filterwarnings('ignore')

WORKER_NAME = sys.argv[1] if len(sys.argv) > 1 else 'worker_test'
TASK_FILE = f'{WORKER_NAME}_tasks.csv'
OUTPUT_FILE = f'{WORKER_NAME}_dividend.csv'
LOG_FILE = f'{WORKER_NAME}.log'

SLEEP_SECONDS = 10  # cninfo 接口间隔
BATCH_SIZE = 50
BATCH_REST_SECONDS = 30


def log(msg):
    line = f'[{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}] {msg}'
    print(line)
    with open(LOG_FILE, 'a', encoding='utf-8') as f:
        f.write(line + '\n')


def parse_amount(x):
    """解析金额/比例字段，处理 None/NaN/逗号/万/亿/%"""
    if pd.isna(x):
        return None
    s = str(x).strip()
    if s in ['', 'None', 'nan', 'NaN', '-']:
        return None
    s = s.replace(',', '').replace('%', '')
    m = re.match(r'^(-?\d+\.?\d*)\s*([万亿])?$', s)
    if not m:
        return None
    num = float(m.group(1))
    unit = m.group(2)
    if unit == '万':
        num *= 10000
    elif unit == '亿':
        num *= 100000000
    return num


def extract_dividend_per_share(row):
    """从 '实施方案分红说明' 中提取每股分红（元）
    例如 '10派10元(含税)' -> 1.0
         '10转增5股派12元(含税)' -> 1.2
    """
    desc = str(row.get('实施方案分红说明', ''))
    if not desc:
        return None
    # 匹配 "派X元"
    m = re.search(r'派(\d+\.?\d*)元', desc)
    if m:
        return float(m.group(1)) / 10
    return None


def extract_stock_dividend(row):
    """提取送转股比例，10送X或10转增X -> 每股送转X/10"""
    desc = str(row.get('实施方案分红说明', ''))
    if not desc:
        return None, None
    send = re.search(r'10送(\d+\.?\d*)股', desc)
    transfer = re.search(r'10转增(\d+\.?\d*)股', desc)
    return (float(send.group(1)) / 10 if send else None,
            float(transfer.group(1)) / 10 if transfer else None)


def fetch_one(code):
    """拉取单只股票分红数据，返回 DataFrame"""
    try:
        time.sleep(SLEEP_SECONDS)
        df = ak.stock_dividend_cninfo(symbol=code)
        if df is None or len(df) == 0:
            return pd.DataFrame()
        return df
    except Exception as e:
        log(f'{code} ERROR: {e}')
        return pd.DataFrame()


def process_and_append(code, name, df, output_path):
    """解析分红数据并追加写入 CSV"""
    if len(df) == 0:
        return 0

    records = []
    for _, row in df.iterrows():
        dividend_per_share = extract_dividend_per_share(row)
        send_ratio, transfer_ratio = extract_stock_dividend(row)
        records.append({
            'code': code,
            'name': name,
            'announce_date': str(row.get('实施方案公告日期')),
            'record_date': str(row.get('股权登记日')),
            'ex_date': str(row.get('除权日')),
            'pay_date': str(row.get('派息日')),
            'dividend_type': str(row.get('分红类型')),
            'dividend_per_share': dividend_per_share,
            'send_ratio': send_ratio,
            'transfer_ratio': transfer_ratio,
            'raw_desc': str(row.get('实施方案分红说明')),
            'report_period': str(row.get('报告时间')),
            'source': 'akshare_cninfo',
            'fetched_at': datetime.now(),
        })

    out_df = pd.DataFrame(records)
    header = not os.path.exists(output_path)
    out_df.to_csv(output_path, mode='a', index=False, header=header)
    return len(records)


def main():
    log(f'Worker {WORKER_NAME} 启动')

    if not os.path.exists(TASK_FILE):
        log(f'任务文件不存在: {TASK_FILE}')
        return

    tasks = pd.read_csv(TASK_FILE)
    log(f'任务数: {len(tasks)}')

    # 如果输出文件已存在，读取已完成的 code
    done_codes = set()
    if os.path.exists(OUTPUT_FILE):
        out_df = pd.read_csv(OUTPUT_FILE)
        done_codes = set(out_df['code'].astype(str).str.zfill(6).tolist())
        log(f'已完成: {len(done_codes)} 只，继续剩余任务')

    total = 0
    for idx, row in tasks.iterrows():
        code = str(row['code']).zfill(6)
        name = row.get('name', '')

        if code in done_codes:
            continue

        log(f'[{idx+1}/{len(tasks)}] 拉取 {code} {name}')
        df = fetch_one(code)
        count = process_and_append(code, name, df, OUTPUT_FILE)
        total += count
        log(f'  -> {count} 条记录')

        if (idx + 1) % BATCH_SIZE == 0:
            log(f'已完成 {idx+1} 只，休息 {BATCH_REST_SECONDS} 秒')
            time.sleep(BATCH_REST_SECONDS)

    log(f'Worker {WORKER_NAME} 完成，总记录: {total}')


if __name__ == '__main__':
    main()