#!/usr/bin/env python3
"""
使用mootdx获取A股历史K线数据
mootdx通过TCP协议直连通达信服务器，不依赖HTTP API，稳定性极高
"""

import sys
import json
from datetime import datetime

def get_kline_data(code: str, days: int = 90):
    """
    获取股票历史K线数据
    
    Args:
        code: 6位股票代码
        days: 需要获取的天数
    
    Returns:
        K线数据列表
    """
    try:
        from mootdx.quotes import Quotes
        
        # 判断市场
        market = 'sh' if code.startswith('6') else 'sz'
        
        # 创建客户端
        client = Quotes.factory(market='std', multithread=True, heartbeat=True)
        
        # 获取K线数据（日K）
        # category: 0=5分钟, 1=15分钟, 2=30分钟, 3=1小时, 4=日K, 5=周K, 6=月K
        klines = client.bars(symbol=code, market=market, category=4, count=days)
        
        if not klines or len(klines) == 0:
            return {
                'error': '未获取到K线数据',
                'data': [],
                'count': 0
            }
        
        # 转换为标准格式
        result = []
        for bar in klines:
            # mootdx返回的是namedtuple，包含: datetime, open, high, low, close, amount, vol
            result.append({
                'date': bar.datetime.strftime('%Y-%m-%d') if hasattr(bar.datetime, 'strftime') else str(bar.datetime)[:10],
                'open': float(bar.open),
                'close': float(bar.close),
                'high': float(bar.high),
                'low': float(bar.low),
                'volume': int(bar.vol) if hasattr(bar, 'vol') else 0,
            })
        
        # 按日期排序（从旧到新）
        result.sort(key=lambda x: x['date'])
        
        return {
            'data': result,
            'count': len(result),
            'source': 'mootdx'
        }
        
    except ImportError:
        return {
            'error': 'mootdx库未安装',
            'data': [],
            'count': 0
        }
    except Exception as e:
        return {
            'error': f'获取失败: {str(e)}',
            'data': [],
            'count': 0
        }

if __name__ == '__main__':
    # 从命令行参数获取股票代码和天数
    if len(sys.argv) < 2:
        print(json.dumps({'error': '请提供股票代码'}))
        sys.exit(1)
    
    code = sys.argv[1]
    days = int(sys.argv[2]) if len(sys.argv) > 2 else 90
    
    result = get_kline_data(code, days)
    print(json.dumps(result, ensure_ascii=False))