# coding:utf-8
"""
同花顺行业资金流向采集工具
接口：moneyflow_ind_ths
文档：https://tushare.pro/document/2?doc_id=343
"""

import os
import sys
from datetime import datetime
from typing import Optional

import pandas as pd
import tushare as ts
from sqlalchemy import create_engine, text

PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(PROJECT_ROOT)

from src.scripts.config import TUSHARE_TOKEN


class IndustryMoneyflowTHSCollector:
    """同花顺行业资金流向采集器"""

    def __init__(self, db_url: str, tushare_token: str, table_name: str = "gp_industry_moneyflow_ths"):
        self.engine = create_engine(
            db_url,
            pool_size=5,
            max_overflow=10,
            pool_recycle=3600,
        )
        self.table_name = table_name

        ts.set_token(tushare_token)
        self.pro = ts.pro_api()

        print("=" * 60)
        print("同花顺行业资金流向采集工具")
        print(f"目标数据表: {self.table_name}")
        print("=" * 60)

    def fetch_data(
        self,
        ts_code: Optional[str] = None,
        trade_date: Optional[str] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
    ) -> pd.DataFrame:
        """调用 Tushare moneyflow_ind_ths 接口"""
        try:
            df = self.pro.moneyflow_ind_ths(
                ts_code=ts_code,
                trade_date=trade_date,
                start_date=start_date,
                end_date=end_date,
            )
            return df
        except Exception as exc:
            print(f"Tushare moneyflow_ind_ths 调用失败: {exc}")
            return pd.DataFrame()

    @staticmethod
    def _to_numeric(series):
        return pd.to_numeric(series, errors="coerce")

    def transform_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """转换接口数据为数据库格式"""
        if df.empty:
            return pd.DataFrame()

        result = pd.DataFrame()
        result["trade_date"] = pd.to_datetime(df["trade_date"], format="%Y%m%d")
        result["industry_code"] = df.get("ts_code")
        result["industry_name"] = df.get("industry")
        result["lead_stock"] = df.get("lead_stock")
        result["close_index"] = self._to_numeric(df.get("close"))
        result["pct_change"] = self._to_numeric(df.get("pct_change"))
        result["company_num"] = self._to_numeric(df.get("company_num")).astype("Int64")
        result["pct_change_stock"] = self._to_numeric(df.get("pct_change_stock"))
        result["close_price"] = self._to_numeric(df.get("close_price"))
        result["net_buy_amount"] = self._to_numeric(df.get("net_buy_amount"))
        result["net_sell_amount"] = self._to_numeric(df.get("net_sell_amount"))
        result["net_amount"] = self._to_numeric(df.get("net_amount"))

        result = result.dropna(subset=["industry_code"])
        result["created_at"] = datetime.now()
        result["updated_at"] = datetime.now()
        return result

    def delete_by_date(self, trade_date: str) -> None:
        """删除指定日期旧数据"""
        with self.engine.begin() as conn:
            delete_sql = text(f"DELETE FROM {self.table_name} WHERE trade_date = :trade_date")
            affected = conn.execute(delete_sql, {"trade_date": trade_date})
            print(f"已删除 {trade_date} 旧数据 {affected.rowcount} 条")

    def save_dataframe(self, df: pd.DataFrame) -> None:
        """写入数据到表"""
        if df.empty:
            print("无数据需要写入")
            return
        df.to_sql(self.table_name, self.engine, if_exists="append", index=False)
        print(f"写入 {len(df)} 条记录")

    def run_daily_collection(self, date: Optional[str] = None) -> None:
        """采集单日数据"""
        target_date = date or datetime.now().strftime("%Y-%m-%d")
        trade_date = datetime.strptime(target_date, "%Y-%m-%d").strftime("%Y%m%d")

        print(f"开始采集 {target_date} ({trade_date}) 行业资金流向数据")
        self.delete_by_date(target_date)

        df = self.fetch_data(trade_date=trade_date)
        if df.empty:
            print("接口未返回数据（可能为非交易日或权限不足）")
            return

        result_df = self.transform_data(df)
        if result_df.empty:
            print("数据转换失败")
            return

        self.save_dataframe(result_df)
        print("每日采集完成")

    def run_range_collection(self, start_date: str, end_date: str) -> None:
        """在日期区间内逐日采集"""
        start = datetime.strptime(start_date, "%Y-%m-%d")
        end = datetime.strptime(end_date, "%Y-%m-%d")
        if start > end:
            raise ValueError("start_date 不能大于 end_date")

        current = start
        while current <= end:
            target = current.strftime("%Y-%m-%d")
            try:
                self.run_daily_collection(target)
            except Exception as exc:
                print(f"{target} 采集失败: {exc}")
            current += pd.Timedelta(days=1)

    def run_full_collection(self, start_date: str = "2010-01-01") -> None:
        """
        全量采集：清空表后，从 start_date (含) 至今逐个交易日采集
        """
        print("=" * 60)
        print("开始执行行业资金流向全量采集")
        print("=" * 60)

        try:
            with self.engine.begin() as conn:
                conn.execute(text(f"TRUNCATE TABLE {self.table_name}"))
            print(f"{self.table_name} 已清空")

            trade_cal = self.pro.trade_cal(
                start_date=start_date.replace("-", ""),
                end_date=datetime.now().strftime("%Y%m%d"),
                fields="cal_date,is_open"
            )
            trade_days = trade_cal.loc[trade_cal["is_open"] == 1, "cal_date"].tolist()
            print(f"共 {len(trade_days)} 个交易日待采集")

            processed = 0
            for cal_date in trade_days:
                target = datetime.strptime(cal_date, "%Y%m%d").strftime("%Y-%m-%d")
                try:
                    self.run_daily_collection(target)
                    processed += 1
                except Exception as exc:
                    print(f"{target} 数据采集失败: {exc}")
                    continue

            print("=" * 60)
            print(f"全量采集完成，成功采集 {processed} 个交易日")
            print("=" * 60)
        except Exception as exc:
            print(f"全量采集失败: {exc}")


def collect_industry_moneyflow_ths(
    db_url: str,
    tushare_token: str,
    mode: str = "daily",
    date: Optional[str] = None,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
) -> None:
    """
    调度入口
    mode:
        - daily: 指定日期（默认今天）数据
        - range: 指定日期区间逐日采集
        - full: 全量覆盖采集（会清空表）
    """
    collector = IndustryMoneyflowTHSCollector(db_url, tushare_token)

    if mode == "daily":
        collector.run_daily_collection(date)
    elif mode == "range":
        if not start_date or not end_date:
            raise ValueError("range 模式需要提供 start_date 和 end_date")
        collector.run_range_collection(start_date, end_date)
    elif mode == "full":
        collector.run_full_collection(start_date or "2010-01-01")
    else:
        raise ValueError(f"未知的采集模式: {mode}")


if __name__ == "__main__":
    # DB_URL = "mysql+pymysql://root:Chlry#$.8@192.168.18.199:3306/db_gp_cj"
    DB_URL = "mysql+pymysql://fac_pattern:Chlry$%.8pattern@192.168.16.150:3306/factordb_mysql"
    TOKEN = TUSHARE_TOKEN

    # 示例：
    collect_industry_moneyflow_ths(DB_URL, TOKEN, mode="daily", date="2025-11-18")
    # collect_industry_moneyflow_ths(DB_URL, TOKEN, mode="range", start_date="2025-11-01", end_date="2025-11-07")
    # collect_industry_moneyflow_ths(DB_URL, TOKEN, mode="full", start_date="2018-01-01")


