#!/usr/bin/env python3
"""Distributed worker: fetch deduced profit for assigned stocks.

Usage on each worker machine:
    export QUANT_WORKER_ID=worker_41
    nohup python3 distributed_worker_deduced_profit.py > worker_41.log 2>&1 &
"""
import os
import sys
import time
import csv
import re
import traceback
import akshare as ak
import psycopg2

WORKER_ID = os.environ.get('QUANT_WORKER_ID')
if not WORKER_ID:
    print("ERROR: set QUANT_WORKER_ID env var, e.g. worker_41")
    sys.exit(1)

OUTPUT_DIR = os.path.expanduser(f"~/elite_stocks_import/{WORKER_ID}")
os.makedirs(OUTPUT_DIR, exist_ok=True)

LOG_FILE = os.path.join(OUTPUT_DIR, f"{WORKER_ID}.log")

PG_CONFIG_106 = {
    'host': '192.168.192.106',
    'port': 5432,
    'dbname': 'quant',
    'user': 'quant_worker',
    'password': 'quant123',
}

DATA_TYPE = 'deduced_profit'
MAX_RETRIES = 3
FETCH_INTERVAL = 2.5  # seconds between akshare requests


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


def parse_amount(s):
    """Parse strings like '412.67亿', '-12.5%', 'False' to float."""
    if s is None or s is False or str(s).strip() == 'False':
        return None
    s = str(s).strip()
    s = s.replace(',', '')
    if s.endswith('亿'):
        return float(s[:-1])
    if s.endswith('万'):
        return float(s[:-1]) / 10000
    if s.endswith('千'):
        return float(s[:-1]) / 100000
    if s.endswith('%'):
        return float(s[:-1])
    try:
        return float(s)
    except:
        return None


def fetch_stock(code):
    """Fetch single stock data from akshare. Returns DataFrame or None."""
    time.sleep(FETCH_INTERVAL)
    df = ak.stock_financial_abstract_ths(symbol=code, indicator='按单季度')
    if df is None or df.empty:
        return None

    column_map = {
        '报告期': 'report_date',
        '净利润': 'net_profit',
        '净利润同比增长率': 'net_profit_yoy',
        '扣非净利润': 'deduced_profit',
        '扣非净利润同比增长率': 'deduced_profit_yoy',
        '营业总收入': 'total_revenue',
        '营业总收入同比增长率': 'total_revenue_yoy',
        '基本每股收益': 'basic_eps',
        '每股净资产': 'bps',
        '净资产收益率-摊薄': 'roe_diluted',
        '销售毛利率': 'gross_margin',
        '销售净利率': 'net_margin',
    }

    available_cols = [c for c in column_map.keys() if c in df.columns]
    selected = df[available_cols].copy()
    selected.rename(columns={c: column_map[c] for c in available_cols}, inplace=True)

    for en_col in column_map.values():
        if en_col not in selected.columns:
            selected[en_col] = None

    selected['code'] = code
    selected['name'] = ''
    selected['source'] = 'akshare_ths'
    selected['fetched_by'] = WORKER_ID
    selected['fetched_at'] = time.strftime('%Y-%m-%d %H:%M:%S')

    numeric_cols = ['net_profit', 'net_profit_yoy', 'deduced_profit', 'deduced_profit_yoy',
                    'total_revenue', 'total_revenue_yoy', 'basic_eps', 'bps',
                    'roe_diluted', 'gross_margin', 'net_margin']
    for col in numeric_cols:
        if col in selected.columns:
            selected[col] = selected[col].apply(parse_amount)

    cols = ['code', 'name', 'report_date', 'net_profit', 'net_profit_yoy',
            'deduced_profit', 'deduced_profit_yoy', 'total_revenue', 'total_revenue_yoy',
            'basic_eps', 'bps', 'roe_diluted', 'gross_margin', 'net_margin',
            'source', 'fetched_by', 'fetched_at']
    selected = selected[cols]
    return selected


def process_one_task(cur, conn, code):
    """Process one pending task. Returns True on success."""
    output_file = os.path.join(OUTPUT_DIR, f"deduced_profit_{code}.csv")

    try:
        log(f"[{code}] pulling...")
        df = fetch_stock(code)
        if df is None or df.empty:
            log(f"[{code}] empty result")
            cur.execute("""
                UPDATE fetch_tasks
                SET status='failed', finished_at=NOW(), error_msg='empty result'
                WHERE code=%s AND data_type=%s
            """, (code, DATA_TYPE))
            conn.commit()
            return False

        df.to_csv(output_file, index=False, encoding='utf-8')
        log(f"[{code}] saved {len(df)} rows to {output_file}")

        cur.execute("""
            UPDATE fetch_tasks
            SET status='done', finished_at=NOW(), output_file=%s, retry_count=retry_count+1
            WHERE code=%s AND data_type=%s
        """, (output_file, code, DATA_TYPE))
        conn.commit()
        return True
    except Exception as e:
        err_msg = str(e)[:500]
        log(f"[{code}] error: {err_msg}")
        log(traceback.format_exc())
        cur.execute("""
            UPDATE fetch_tasks
            SET status='failed', finished_at=NOW(), error_msg=%s, retry_count=retry_count+1
            WHERE code=%s AND data_type=%s
        """, (err_msg, code, DATA_TYPE))
        conn.commit()
        return False


def get_pending_tasks(cur):
    """Fetch pending tasks assigned to this worker, ordered by retry_count."""
    cur.execute("""
        SELECT code, retry_count FROM fetch_tasks
        WHERE data_type = %s AND worker_id = %s AND status = 'pending'
        ORDER BY retry_count ASC, id ASC
        LIMIT 100
    """, (DATA_TYPE, WORKER_ID))
    return cur.fetchall()


def main():
    log(f"Worker {WORKER_ID} starting, output dir: {OUTPUT_DIR}")

    conn = psycopg2.connect(**PG_CONFIG_106)
    conn.autocommit = False
    cur = conn.cursor()

    total = 0
    success = 0
    failed = 0

    while True:
        tasks = get_pending_tasks(cur)
        if not tasks:
            log("No pending tasks, worker finished.")
            break

        for code, retry_count in tasks:
            if retry_count >= MAX_RETRIES:
                log(f"[{code}] exceeded max retries, skipping")
                cur.execute("""
                    UPDATE fetch_tasks
                    SET status='failed', finished_at=NOW(), error_msg='exceeded max retries'
                    WHERE code=%s AND data_type=%s
                """, (code, DATA_TYPE))
                conn.commit()
                failed += 1
                continue

            # Mark as started
            cur.execute("""
                UPDATE fetch_tasks
                SET status='running', started_at=NOW()
                WHERE code=%s AND data_type=%s
            """, (code, DATA_TYPE))
            conn.commit()

            total += 1
            if process_one_task(cur, conn, code):
                success += 1
            else:
                failed += 1

    log(f"Worker {WORKER_ID} finished: total={total}, success={success}, failed={failed}")
    conn.close()


if __name__ == '__main__':
    main()