"""
封装adata库的A股数据获取方法
提供简单易用的接口来获取股票市场数据
"""

import adata
from typing import Dict, List, Optional
import pandas as pd
from datetime import datetime, timedelta

class AStockAPI:
    @staticmethod
    def get_all_stock_codes() -> List[str]:
        """
        获取所有A股代码
        Returns:
            List[str]: 股票代码列表
        """
        try:
            return adata.stock.info.all_code()
        except Exception as e:
            print(f"获取股票代码失败: {str(e)}")
            return []


    @staticmethod
    def get_mine_clearance_tdx(stock_code: str) -> pd.DataFrame:
        """
        获取股票扫雷避险信息
        """
        try:
            return adata.sentiment.mine.mine_clearance_tdx(stock_code=stock_code)
        except Exception as e:
            print(f"获取情绪数据失败: {str(e)}")
            return []

    @staticmethod
    def get_concept_east(stock_code: str) -> pd.DataFrame:
        """
        获取股票所属的概念板块
        """
        try:
            return adata.stock.info.get_concept_east(stock_code)
        except Exception as e:
            print(f"获取概念板块失败: {str(e)}")
            return pd.DataFrame()


    @staticmethod
    def get_plate_east(stock_code: str) -> pd.DataFrame:
        """
        获取股票所属的板块
        """
        try:
            return adata.stock.info.get_plate_east(stock_code)
        except Exception as e:
            print(f"获取板块失败: {str(e)}")
            return pd.DataFrame()

    @staticmethod
    def get_stock_shares(stock_code: str) -> pd.DataFrame:
        """
        获取股本信息
        """
        try:
            return adata.stock.info.get_stock_shares(stock_code)
        except Exception as e:
            print(f"获取股本信息失败: {str(e)}")
            return pd.DataFrame()

    @staticmethod
    def get_market_data(stock_code: str, 
                       start_date: Optional[str] = None,
                       end_date: Optional[str] = None,
                       k_type: int = 1,
                       adjust_type: int = 1) -> pd.DataFrame:
        """
        获取股票市场日线数据
        Args:
            codes: 股票代码列表，默认为None表示获取所有股票
            start_date: 开始日期，格式：YYYY-MM-DD
            end_date: 结束日期，格式：YYYY-MM-DD
        Returns:
            pd.DataFrame: 市场数据
        """
        try:
            return adata.stock.market.get_market(stock_code=stock_code, 
                                               start_date=start_date,
                                               end_date=end_date, 
                                               k_type=k_type,
                                               adjust_type=adjust_type)
        except Exception as e:
            print(f"获取市场数据失败: {str(e)}")
            return pd.DataFrame()

    @staticmethod
    def get_market_min_data(stock_code: str) -> pd.DataFrame:
        """
        获取股票市场分钟线数据
        Args:
            stock_code: 股票代码
        Returns:
            pd.DataFrame: 分钟级市场数据
        """
        try:
            return adata.stock.market.get_market_min(stock_code=stock_code)
        except Exception as e:
            print(f"获取分钟线数据失败: {str(e)}")
            return pd.DataFrame()
        
    @staticmethod
    def get_capital_flow(stock_code: str, start_date: Optional[str] = None, end_date: Optional[str] = None) -> pd.DataFrame:
        """
        获取股票资金流向数据
        """
        try:
            return adata.stock.market.get_capital_flow(stock_code=stock_code, start_date=start_date, end_date=end_date)
        except Exception as e:
            print(f"获取资金流向数据失败: {str(e)}")
            return pd.DataFrame()
        
    @staticmethod
    def get_capital_flow_min(stock_code: str) -> pd.DataFrame:
        """
        获取股票资金流向分钟数据
        Args:
            stock_code: 股票代码
        Returns:
            pd.DataFrame: 资金流向数据
        """
        try:
            return adata.stock.market.get_capital_flow_min(stock_code=stock_code)
        except Exception as e:
            print(f"获取资金流向数据失败: {str(e)}")
            return pd.DataFrame()

# 使用示例
if __name__ == "__main__":

    print("开始获取A股数据")
    api = AStockAPI()

    stock_codes = api.get_all_stock_codes()
    list = stock_codes.to_json(orient="records")
    print(list)

    # 保存到数据库
    import cryptoai.utils.db_manager as db_manager
    db_manager.get_db_manager().create_stocks(list)



    # 获取所有股票代码
    
    # # 获取所有股票代码
    # stock_codes = ["688552", "688648", "688165","688552"]
    # print(f"获取到 {len(stock_codes)} 个股票代码")
    
    # # 获取指定股票的市场数据
    # if stock_codes:
    #     sample_codes = stock_codes[:3]  # 取前3个股票作为示例
    #     print(f"获取到 {len(sample_codes)} 个股票代码")
    #     today = datetime.now().strftime("%Y-%m-%d")
    #     week_ago = (datetime.now() - timedelta(days=7)).strftime("%Y-%m-%d")
        
    #     # 获取日线数据
    #     market_data = api.get_market_data(stock_code=sample_codes[0],
    #                                     start_date=week_ago,
    #                                     end_date=today)
    #     print("\n日线数据示例:")
    #     print(market_data.head())
        
    #     # 获取分钟线数据
    #     min_data = api.get_market_min_data(stock_code=sample_codes[0])
    #     print("\n分钟线数据示例:")
    #     print(min_data.head())
        
    #     # 获取资金流向数据
    #     flow_data = api.get_capital_flow_min(stock_code=sample_codes[0])
    #     print("\n资金流向数据示例:")
    #     print(flow_data.head())