"""
Wind 数据获取模块
提供统一的接口获取股票历史数据、财务数据和技术指标
"""

from WindPy import w
import pandas as pd
from datetime import datetime, timedelta
from typing import List, Dict, Optional, Tuple


class WindDataFetcher:
    """Wind数据获取器"""

    def __init__(self):
        """初始化Wind连接"""
        self._start()

    def _start(self):
        """启动Wind连接"""
        result = w.start()
        if result.ErrorCode != 0:
            raise ConnectionError(f"Wind连接失败: {result.Data}")
        print("Wind连接成功")

    def _stop(self):
        """关闭Wind连接"""
        w.stop()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._stop()

    @staticmethod
    def _handle_wsd_result(result, field_names):
        """处理wsd返回结果"""
        if result.ErrorCode != 0:
            raise ValueError(f"数据获取失败: {result.Data}")

        # 转换为DataFrame
        data_dict = dict(zip(field_names, result.Data))
        df = pd.DataFrame(data_dict, columns=field_names)
        df['date'] = result.Times
        return df

    def get_daily_data(self,
                      stock_code: str,
                      start_date: str,
                      end_date: str,
                      fields: List[str]) -> pd.DataFrame:
        """
        获取日线数据

        Args:
            stock_code: 股票代码，如 '300866.SZ'
            start_date: 开始日期，格式 'YYYY-MM-DD'
            end_date: 结束日期，格式 'YYYY-MM-DD'
            fields: 数据字段列表

        Returns:
            DataFrame包含日期和指定字段的数据
        """
        result = w.wsd(stock_code, fields, start_date, end_date)
        return self._handle_wsd_result(result, fields)

    def get_stock_basic(self, stock_code: str) -> Dict:
        """
        获取股票基本信息

        Args:
            stock_code: 股票代码

        Returns:
            包含股票名称、行业、上市日期等信息的字典
        """
        fields = ['sec_name', 'industry', 'ipo_date', 'list_date', 'trade_status']
        result = w.wss(stock_code, fields)

        if result.ErrorCode != 0:
            raise ValueError(f"基本信息获取失败: {result.Data}")

        return dict(zip(fields, result.Data[0]))

    def get_technical_indicators(self,
                                stock_code: str,
                                start_date: str,
                                end_date: str) -> pd.DataFrame:
        """
        获取技术指标数据

        Args:
            stock_code: 股票代码
            start_date: 开始日期
            end_date: 结束日期

        Returns:
            DataFrame包含MA、成交量等技术指标
        """
        fields = [
            'close', 'open', 'high', 'low', 'volume', 'amt', 'pct_chg',
            'ma5', 'ma10', 'ma20', 'ma60',
            'pe_ttm', 'pb_lf', 'ps_ttm', 'pcf_ncf_ttm'
        ]
        return self.get_daily_data(stock_code, start_date, end_date, fields)

    def get_financial_data(self,
                          stock_code: str,
                          report_date: str = '') -> Dict:
        """
        获取财务数据

        Args:
            stock_code: 股票代码
            report_date: 报告期，格式 'YYYYMMDD'，默认最新

        Returns:
            包含ROE、营收、净利润等财务指标的字典
        """
        fields = [
            'roe_wgt', 'roa2', 'net_profit_to_profit', 'total_revenue_ps',
            'profit_to_gr', 'ebit_ps', 'assets_to_eqt', 'debt_to_assets',
            'current_ratio', 'quick_ratio', 'op_income_to_revenue'
        ]

        result = w.wss(stock_code, fields, 'rptDate={}'.format(report_date))

        if result.ErrorCode != 0:
            raise ValueError(f"财务数据获取失败: {result.Data}")

        return dict(zip(fields, result.Data[0]))

    def get_weekly_data(self,
                       stock_code: str,
                       start_date: str,
                       end_date: str,
                       fields: List[str] = None) -> pd.DataFrame:
        """
        获取并汇总周线数据

        Args:
            stock_code: 股票代码
            start_date: 开始日期
            end_date: 结束日期
            fields: 数据字段，默认为基本OHLCV

        Returns:
            DataFrame包含周线数据
        """
        if fields is None:
            fields = ['close', 'open', 'high', 'low', 'volume', 'amt', 'pct_chg']

        # 获取日线数据
        df_daily = self.get_daily_data(stock_code, start_date, end_date, fields)

        # 转换为周线数据
        df_daily.set_index('date', inplace=True)
        df_weekly = df_daily.resample('W').agg({
            'open': 'first',
            'high': 'max',
            'low': 'min',
            'close': 'last',
            'volume': 'sum',
            'amt': 'sum',
            'pct_chg': 'sum'  # 周涨跌幅近似为日涨跌幅之和
        }).dropna()

        df_weekly.reset_index(inplace=True)
        return df_weekly

    def get_stock_list_by_sector(self, sector_name: str, limit: int = 10) -> List[str]:
        """
        获取某行业的龙头股票列表

        Args:
            sector_name: 行业名称
            limit: 返回数量

        Returns:
            股票代码列表
        """
        # 使用Wind的行业板块数据
        fields = ['sec_name', 'ipo_date']
        # 这里简化处理，实际需要根据Wind板块接口调整
        print(f"注意: 行业股票列表获取需要根据实际Wind板块接口实现")
        return []

    def calculate_volatility(self, df: pd.DataFrame, annualized: bool = True) -> float:
        """
        计算波动率

        Args:
            df: 包含收盘价的DataFrame
            annualized: 是否年化

        Returns:
            波动率（百分比）
        """
        returns = df['close'].pct_change().dropna()
        volatility = returns.std()

        if annualized:
            # 日数据年化系数：sqrt(252)，周数据年化系数：sqrt(52)
            n = len(df)
            if n > 250:  # 日线数据
                volatility *= 252 ** 0.5
            else:  # 周线数据
                volatility *= 52 ** 0.5

        return volatility * 100  # 转换为百分比


# 便捷函数
def fetch_stock_report_data(stock_code: str,
                           days_back: int = 90,
                           weekly: bool = True) -> Tuple[pd.DataFrame, Dict]:
    """
    获取股票报告所需的数据

    Args:
        stock_code: 股票代码
        days_back: 获取多少天的数据
        weekly: 是否为周线数据

    Returns:
        (price_df, financial_dict) 价格数据和财务数据
    """
    end_date = datetime.now().strftime('%Y-%m-%d')
    start_date = (datetime.now() - timedelta(days=days_back)).strftime('%Y-%m-%d')

    with WindDataFetcher() as fetcher:
        # 获取价格数据
        if weekly:
            price_df = fetcher.get_weekly_data(stock_code, start_date, end_date)
        else:
            price_df = fetcher.get_technical_indicators(stock_code, start_date, end_date)

        # 获取财务数据
        financial_data = fetcher.get_financial_data(stock_code)

        # 获取基本信息
        basic_info = fetcher.get_stock_basic(stock_code)

    # 合并财务信息
    financial_data.update(basic_info)

    return price_df, financial_data


if __name__ == '__main__':
    # 测试代码
    with WindDataFetcher() as fetcher:
        # 示例：获取某股票的数据
        code = '300866.SZ'
        start = '2025-01-01'
        end = '2026-02-06'

        print(f"获取 {code} 的日线数据...")
        df = fetcher.get_technical_indicators(code, start, end)
        print(df.tail())

        print(f"\n获取 {code} 的基本信息...")
        info = fetcher.get_stock_basic(code)
        print(info)