"""
市场概览查询

用法: python market_overview.py [--no-indices] [--no-sectors] [--no-hk]
  --no-indices: 不包含A股指数数据
  --no-sectors: 不包含板块数据
  --no-hk:      不包含港股指数数据

AKShare接口:
  A股指数:
    主: ak.stock_zh_index_spot_sina() — 新浪A股指数行情（稳定）
    备: ak.stock_zh_index_spot_em() — 东方财富A股指数行情
  港股指数:
    主: ak.stock_hk_index_spot_sina() — 新浪港股指数行情（稳定）
    备: ak.stock_hk_index_spot_em() — 东方财富港股指数行情
  板块:
    主: ak.stock_board_industry_name_em() — 东方财富行业板块

注意: 新浪接口频繁调用会被封IP
"""

import sys
import argparse

sys.path.insert(0, '.')
from _utils import output_json, output_error, check_akshare, safe_float


# A股主要指数代码
MAIN_INDICES = {
    'sh000001': '上证指数',
    'sz399001': '深证成指',
    'sz399006': '创业板指',
    'sh000300': '沪深300',
    'sh000016': '上证50',
    'sz399005': '中小板指',
}

# 港股主要指数代码（新浪代码格式）
HK_MAIN_INDICES = {
    'HSI': '恒生指数',
    'HSTECH': '恒生科技指数',
    'HSCEI': '恒生中国企业指数',
    'HSCCI': '恒生香港中资企业指数(红筹)',
}


def _fetch_indices_em():
    """东方财富A股指数"""
    import akshare as ak

    indices = []
    try:
        index_df = ak.stock_zh_index_spot_em(symbol="上证系列指数")
        key_codes = {'000001', '000016', '000300'}
        for _, row in index_df.iterrows():
            idx_code = str(row.get('代码', ''))
            if idx_code in key_codes:
                indices.append({
                    "code": idx_code,
                    "name": str(row.get('名称', '')),
                    "price": safe_float(row.get('最新价')),
                    "change_pct": safe_float(row.get('涨跌幅')),
                    "change_amt": safe_float(row.get('涨跌额')),
                    "volume": safe_float(row.get('成交量')),
                    "amount": safe_float(row.get('成交额')),
                })
    except Exception:
        pass

    try:
        sz_df = ak.stock_zh_index_spot_em(symbol="深证系列指数")
        sz_key = {'399001', '399006', '399005'}
        for _, row in sz_df.iterrows():
            idx_code = str(row.get('代码', ''))
            if idx_code in sz_key:
                indices.append({
                    "code": idx_code,
                    "name": str(row.get('名称', '')),
                    "price": safe_float(row.get('最新价')),
                    "change_pct": safe_float(row.get('涨跌幅')),
                    "change_amt": safe_float(row.get('涨跌额')),
                    "volume": safe_float(row.get('成交量')),
                    "amount": safe_float(row.get('成交额')),
                })
    except Exception:
        pass

    return indices


def _fetch_indices_sina():
    """新浪A股指数（主接口）"""
    import akshare as ak

    df = ak.stock_zh_index_spot_sina()

    indices = []
    for _, row in df.iterrows():
        idx_code = str(row.get('代码', ''))
        if idx_code in MAIN_INDICES:
            indices.append({
                "code": idx_code,
                "name": str(row.get('名称', MAIN_INDICES.get(idx_code, ''))),
                "price": safe_float(row.get('最新价')),
                "change_pct": safe_float(row.get('涨跌幅')),
                "change_amt": safe_float(row.get('涨跌额')),
                "volume": safe_float(row.get('成交量')),
                "amount": safe_float(row.get('成交额')),
                "data_source": "sina",
            })

    return indices


def _fetch_hk_indices_sina():
    """新浪港股指数"""
    import akshare as ak

    df = ak.stock_hk_index_spot_sina()

    indices = []
    for _, row in df.iterrows():
        idx_code = str(row.get('代码', ''))
        idx_name = str(row.get('名称', ''))
        # 匹配主要港股指数
        if idx_code in HK_MAIN_INDICES or any(kw in idx_name for kw in ['恒生指数', '恒生科技', '国企指数', '红筹']):
            indices.append({
                "code": idx_code,
                "name": idx_name,
                "price": safe_float(row.get('最新价')),
                "change_pct": safe_float(row.get('涨跌幅')),
                "change_amt": safe_float(row.get('涨跌额')),
                "data_source": "sina_hk",
            })

    # 如果没匹配到关键词，取前5个
    if not indices and not df.empty:
        for _, row in df.head(5).iterrows():
            indices.append({
                "code": str(row.get('代码', '')),
                "name": str(row.get('名称', '')),
                "price": safe_float(row.get('最新价')),
                "change_pct": safe_float(row.get('涨跌幅')),
                "change_amt": safe_float(row.get('涨跌额')),
                "data_source": "sina_hk",
            })

    return indices


def _fetch_sectors_em():
    """东方财富行业板块"""
    import akshare as ak

    sector_df = ak.stock_board_industry_name_em()
    sector_df = sector_df.sort_values(by='涨跌幅', ascending=False)

    top_up = []
    for _, row in sector_df.head(5).iterrows():
        top_up.append({
            "name": str(row.get('板块名称', '')),
            "change_pct": safe_float(row.get('涨跌幅')),
            "leading_stock": str(row.get('领涨股票', '')),
            "leading_pct": safe_float(row.get('领涨股票涨跌幅')),
        })

    top_down = []
    for _, row in sector_df.tail(5).iterrows():
        top_down.append({
            "name": str(row.get('板块名称', '')),
            "change_pct": safe_float(row.get('涨跌幅')),
            "leading_stock": str(row.get('领跌股票', '')),
            "leading_pct": safe_float(row.get('领跌股票涨跌幅')),
        })

    return top_up, top_down, len(sector_df)


def main():
    parser = argparse.ArgumentParser(description='市场概览查询')
    parser.add_argument('--no-indices', action='store_true', help='不包含A股指数数据')
    parser.add_argument('--no-sectors', action='store_true', help='不包含板块数据')
    parser.add_argument('--no-hk', action='store_true', help='不包含港股指数数据')
    args = parser.parse_args()

    err = check_akshare()
    if err:
        output_error(err)
        return

    include_indices = not args.no_indices
    include_sectors = not args.no_sectors
    include_hk = not args.no_hk

    try:
        data = {}

        # A股主要指数：优先新浪，失败回退东方财富
        if include_indices:
            indices = []
            try:
                indices = _fetch_indices_sina()
            except Exception:
                pass

            if not indices:
                try:
                    indices = _fetch_indices_em()
                except Exception as e:
                    data["indices_error"] = str(e)

            data["indices"] = indices

        # 港股指数
        if include_hk:
            hk_indices = []
            try:
                hk_indices = _fetch_hk_indices_sina()
            except Exception as e:
                data["hk_indices_error"] = str(e)

            data["hk_indices"] = hk_indices

        # 行业板块涨跌
        if include_sectors:
            try:
                top_up, top_down, total = _fetch_sectors_em()
                data["top_sectors_up"] = top_up
                data["top_sectors_down"] = top_down
                data["total_sectors"] = total
            except Exception as e:
                data["top_sectors_up"] = []
                data["top_sectors_down"] = []
                data["sectors_error"] = str(e)

        output_json(data)

    except Exception as e:
        output_error(f"获取市场概览失败: {str(e)}")


if __name__ == '__main__':
    main()
