#!/usr/bin/env python3
"""
股票实时价格查询工具
支持A股、港股、美股等多个市场的股票价格查询
"""

import sys
import json
import time
from typing import List, Dict, Optional
import yfinance as yf
import requests
from datetime import datetime


class StockPriceQuery:
    """股票价格查询类"""

    def __init__(self):
        self.data_sources = {
            "yahoo_finance": self._query_yahoo_finance,
            "eastmoney": self._query_eastmoney,
            "xueqiu": self._query_xueqiu,
        }
        self.exchange_mapping = {"HK": "HK", "SH": "SS", "SZ": "SZ", "US": ""}

    def standardize_code(self, stock_code: str) -> str:
        """标准化股票代码"""
        stock_code = stock_code.strip()

        # 处理已有后缀的情况
        if "." in stock_code:
            return stock_code

        # 根据代码长度和前缀判断市场
        if len(stock_code) == 5 and stock_code.isdigit():
            # 港股5位数字
            return f"{stock_code}.HK"
        elif len(stock_code) == 6 and stock_code.isdigit():
            # A股6位数字
            if stock_code.startswith(("00", "30")):
                return f"{stock_code}.SZ"  # 深圳
            else:
                return f"{stock_code}.SS"  # 上海
        elif stock_code.isalpha():
            # 美股代码
            return f"{stock_code}"
        else:
            # 默认按港股处理
            return f"{stock_code}.HK"

    def _query_yahoo_finance(self, standardized_code: str) -> Optional[Dict]:
        """查询Yahoo Finance数据"""
        try:
            ticker = yf.Ticker(standardized_code)
            info = ticker.info

            if "currentPrice" not in info:
                return None

            price = info["currentPrice"]
            previous_close = info.get("previousClose", price)
            currency = info.get("currency", "USD")

            return {
                "price": float(price),
                "previous_close": float(previous_close),
                "currency": currency,
                "timestamp": datetime.now().isoformat(),
                "source": "Yahoo Finance",
            }
        except Exception as e:
            print(
                f"Yahoo Finance query failed for {standardized_code}: {e}",
                file=sys.stderr,
            )
            return None

    def _query_eastmoney(self, standardized_code: str) -> Optional[Dict]:
        """查询东方财富数据（简化版）"""
        # 这里可以实现东方财富的API调用
        # 暂时返回None，主要依赖Yahoo Finance
        return None

    def _query_xueqiu(self, standardized_code: str) -> Optional[Dict]:
        """查询雪球数据（简化版）"""
        # 这里可以实现雪球的API调用
        # 暂时返回None，主要依赖Yahoo Finance
        return None

    def get_price(self, stock_code: str) -> Dict:
        """获取股票价格"""
        standardized_code = self.standardize_code(stock_code)
        results = []

        # 尝试多个数据源
        for source_name, query_func in self.data_sources.items():
            try:
                result = query_func(standardized_code)
                if result:
                    result["code"] = stock_code
                    result["standardized_code"] = standardized_code
                    results.append(result)

                    # 如果Yahoo Finance成功，通常就足够了
                    if source_name == "yahoo_finance":
                        break

            except Exception as e:
                print(f"Query {source_name} failed: {e}", file=sys.stderr)
                continue

        if not results:
            return {
                "code": stock_code,
                "standardized_code": standardized_code,
                "error": "无法获取准确价格数据",
                "timestamp": datetime.now().isoformat(),
            }

        # 返回第一个成功的结果
        return results[0]

    def get_prices(self, stock_codes: List[str]) -> List[Dict]:
        """批量获取股票价格"""
        results = []
        for code in stock_codes:
            result = self.get_price(code)
            results.append(result)
            # 避免请求过于频繁
            time.sleep(0.1)
        return results


def main():
    """主函数"""
    if len(sys.argv) < 2:
        print("用法:")
        print("  python stock_price_query.py <stock_code>")
        print("  python stock_price_query.py <stock_code1> <stock_code2> ...")
        print("")
        print("示例:")
        print("  python stock_price_query.py 00700.HK")
        print("  python stock_price_query.py 00700.HK 09868.HK 001309.SZ")
        sys.exit(1)

    stock_codes = sys.argv[1:]

    if len(stock_codes) == 1:
        query = StockPriceQuery()
        result = query.get_price(stock_codes[0])
        print(json.dumps(result, indent=2, ensure_ascii=False))
    else:
        query = StockPriceQuery()
        results = query.get_prices(stock_codes)
        print(json.dumps(results, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()