"""
测试Baostock股票代码格式转换
"""

import sys
import os
import asyncio
import logging
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.data.data_initializer import DataInitializer
from src.config.settings import Settings

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def get_baostock_format_code(stock_code: str) -> str:
	"""
	将股票代码转换为Baostock格式
	"""
	if stock_code.startswith("sh.") or stock_code.startswith("sz."):
		return stock_code
	
	if stock_code.startswith("6"):
		return f"sh.{stock_code}"
	elif stock_code.startswith("0") or stock_code.startswith("3"):
		return f"sz.{stock_code}"
	else:
		return stock_code


async def test_baostock_format():
	"""
	测试Baostock格式股票代码
	"""
	try:
		logger.info("开始测试Baostock格式股票代码...")
		
		# 加载配置
		settings = Settings()
		logger.info("配置加载成功")
		
		# 创建数据初始化器
		initializer = DataInitializer(settings)
		logger.info("数据初始化器创建成功")
		
		# 测试股票代码列表
		test_codes = [
			"000001",  # 平安银行
			"600000",  # 浦发银行
			"300001",  # 特锐德
			"000007"   # 全新好
		]
		
		logger.info("测试股票代码格式转换:")
		for code in test_codes:
			baostock_code = get_baostock_format_code(code)
			logger.info(f"  {code} -> {baostock_code}")
		
		# 测试获取K线数据
		total_kline_data = []
		success_count = 0
		error_count = 0
		
		for code in test_codes:
			try:
				baostock_code = get_baostock_format_code(code)
				logger.info(f"测试获取股票{code}({baostock_code})的K线数据...")
				
				# 使用数据管理器获取K线数据
				kline_data = await initializer.data_manager.get_daily_kline_data(
					baostock_code,
					"2024-01-01",
					"2024-01-10"
				)
				
				if kline_data:
					total_kline_data.extend(kline_data)
					success_count += 1
					logger.info(f"股票{code}获取到{len(kline_data)}条K线数据")
					
					# 打印前几条数据
					for i, data in enumerate(kline_data[:3]):
						logger.info(f"  第{i+1}条: {data}")
				else:
					logger.warning(f"股票{code}未获取到K线数据")
					error_count += 1
				
			except Exception as e:
				logger.error(f"获取股票{code}K线数据失败: {str(e)}")
				error_count += 1
				continue
		
		logger.info(f"测试完成: 成功{success_count}只, 失败{error_count}只, 共获取{len(total_kline_data)}条数据")
		
		return {
			"success": True,
			"test_codes": test_codes,
			"success_count": success_count,
			"error_count": error_count,
			"kline_data_count": len(total_kline_data)
		}
		
	except Exception as e:
		logger.error(f"测试Baostock格式异常: {str(e)}")
		return {"success": False, "error": str(e)}


async def main():
	"""
	主函数
	"""
	result = await test_baostock_format()
	
	if result["success"]:
		logger.info("Baostock格式测试成功！")
		print(f"测试结果: {result}")
		
		if result["kline_data_count"] > 0:
			print("✓ Baostock格式转换成功，可以正常获取K线数据")
		else:
			print("⚠ Baostock格式转换成功，但未获取到K线数据")
	else:
		logger.error("Baostock格式测试失败！")
		print(f"测试失败: {result.get('error')}")
	
	return result


if __name__ == "__main__":
	# 运行测试
	result = asyncio.run(main())
	
	# 输出最终结果
	if result.get("success", False):
		print("\n测试完成！")
		sys.exit(0)
	else:
		print("\n测试失败！")
		sys.exit(1)