"""
快速K线数据更新脚本
只更新部分股票进行测试
"""

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
from src.storage.database import db_manager
from src.storage.stock_repository import StockRepository

# 配置日志
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 update_kline_data_fast():
	"""
	快速更新K线数据（只处理前20只股票）
	"""
	try:
		logger.info("开始快速更新K线数据...")
		
		# 加载配置
		settings = Settings()
		logger.info("配置加载成功")
		
		# 创建数据初始化器
		initializer = DataInitializer(settings)
		logger.info("数据初始化器创建成功")
		
		# 创建存储库
		repository = StockRepository(db_manager.get_session())
		logger.info("存储库创建成功")
		
		# 获取所有股票基础信息
		stocks = repository.get_stock_basic_info()
		logger.info(f"找到{len(stocks)}只股票")
		
		if not stocks:
			logger.error("没有股票基础信息，无法更新")
			return {"success": False, "error": "没有股票基础信息"}
		
		# 只处理前20只股票进行测试
		test_stocks = stocks[:20]
		logger.info(f"测试更新前{len(test_stocks)}只股票")
		
		# 分批处理，每次处理5只股票
		batch_size = 5
		total_batches = (len(test_stocks) + batch_size - 1) // batch_size
		
		total_kline_data = []
		success_count = 0
		error_count = 0
		
		for batch_num in range(total_batches):
			start_idx = batch_num * batch_size
			end_idx = min(start_idx + batch_size, len(test_stocks))
			batch_stocks = test_stocks[start_idx:end_idx]
			
			logger.info(f"处理第{batch_num + 1}批股票，共{len(batch_stocks)}只")
			
			for stock in batch_stocks:
				try:
					# 转换为Baostock格式
					baostock_code = get_baostock_format_code(stock.code)
					logger.info(f"获取股票{stock.code}({baostock_code})的K线数据...")
					
					# 使用数据管理器获取K线数据
					kline_data = await initializer.data_manager.get_daily_kline_data(
						baostock_code,
						"2024-01-01",
						"2024-01-31"  # 获取一个月的数据
					)
					
					if kline_data:
						total_kline_data.extend(kline_data)
						success_count += 1
						logger.info(f"股票{stock.code}获取到{len(kline_data)}条K线数据")
						
						# 保存到数据库
						save_result = repository.save_daily_kline_data(kline_data)
						logger.info(f"股票{stock.code}K线数据保存结果: {save_result}")
					else:
						logger.warning(f"股票{stock.code}未获取到K线数据")
						error_count += 1
					
				except Exception as e:
					logger.error(f"获取股票{stock.code}K线数据失败: {str(e)}")
					error_count += 1
					continue
			
			# 每批处理完成后暂停一下
			await asyncio.sleep(1)
		
		logger.info(f"K线数据更新完成: 成功{success_count}只, 失败{error_count}只, 共获取{len(total_kline_data)}条数据")
		
		return {
			"success": True,
			"test_stocks": len(test_stocks),
			"success_count": success_count,
			"error_count": error_count,
			"kline_data_count": len(total_kline_data)
		}
		
	except Exception as e:
		logger.error(f"更新K线数据异常: {str(e)}")
		return {"success": False, "error": str(e)}


async def main():
	"""
	主函数
	"""
	result = await update_kline_data_fast()
	
	if result["success"]:
		logger.info("K线数据更新成功！")
		print(f"更新结果: {result}")
		
		if result["kline_data_count"] > 0:
			print("✓ K线数据更新成功，数据已保存到数据库")
		else:
			print("⚠ K线数据更新完成，但未获取到数据")
	else:
		logger.error("K线数据更新失败！")
		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)