from datetime import timedelta

from DB.model.StockDailyFreq import get_stock_daily_freq
from DB.sqlite_db_main import SqliteDbMain, config
from utils.tdxUtil import TdxUtil
from utils.comm import *
from 基本信息入库 import StockInfoMain


class StockDailyFreqMain:

	def __init__(self):
		self.code_res = StockInfoMain().get_stock_basic()
		self.db_main = SqliteDbMain(config.stock_daily_freq_db)

	def __filter_stock(self, code, name):
		if 'ST' in name:
			return False
		if code.startswith('30'):
			return False
		if code.startswith('68'):
			return False
		return True

	def init_data(self):
		"""
		全量入库操作
		:return:
		"""
		# Print the results
		for result in self.code_res:
			tdx_util = TdxUtil("")
			s_type = tdx_util.get_security_type(code=result.ts_code)
			if s_type in tdx_util.SECURITY_TYPE and self.__filter_stock(result.ts_code, result.name):
				# 设置开始和结束时间
				start_time = f"{(datetime.now() - timedelta(days=365 * 3)).strftime('%Y-%m-%d')} 09:30:00"
				end_time = f"{datetime.now().strftime('%Y-%m-%d')} 15:00:00"
				print(
					f"{result.id}，{result.ts_code} --> start_time={trade_date} 09:30:00 --->> end_time={trade_date} 15:00:00")
				i = 0
				df = None
				while True:
					try:
						if i > 6:
							break
						df = xcsc_pro.stk_mins(ts_code=result.ts_code, freq='30min',
											   start_time=start_time,
											   end_time=end_time)[::-1]
					except Exception as e:
						i += 1
						# 捕获超时异常
						print("请求超时，等待2分钟后重试...")
						time.sleep(120)  # 休眠2分钟
				print(f"{result.ts_code}，{result.name} 获取 stk_mins _ 30 数据完毕！")
				# 创建表
				table_name = str(result.ts_code).split(".")[0] + "_daily_freq_30"
				new_table_class = get_stock_daily_freq(table_name=table_name)
				self.db_main.create_table(new_table_class)
				entries = []
				for index, row in df.iterrows():
					min_time = row['trade_time']
					# 将字符串解析为日期时间对象
					trade_time = datetime.strptime(min_time, '%Y-%m-%d %H:%M:%S').strftime('%Y-%m-%d')
					entry = new_table_class(
						trade_date=trade_time,
						time=min_time,
						open=row['open'],
						close=row['close'],
						high=row['high'],
						low=row['low'],
						vol=row['vol'],
						amount=row['amount'],
					)
					print(entry.to_dict())
					entries.append(entry)
				self.db_main.insert_all_entry(entries)

	def task_data(self, trade_date=datetime.now().strftime('%Y%m%d')):
		print(len(self.code_res))
		for result in self.code_res:
			tdx_util = TdxUtil("")
			s_type = tdx_util.get_security_type(code=result.ts_code)
			if s_type in tdx_util.SECURITY_TYPE and self.__filter_stock(result.ts_code, result.name):
				print(
					f"{result.id}，{result.ts_code} --> start_time={trade_date} 09:30:00 --->> end_time={trade_date} 15:00:00")

				df = xcsc_pro.stk_mins(ts_code=result.ts_code, freq='30min',
									   start_time=f"{trade_date} 09:30:00",
									   end_time=f"{trade_date} 15:00:00")
				print(f"{result.ts_code},{result.name} 获取 stk_mins _ 30 数据完毕！")
				for index, row in df.iterrows():
					table_name = str(result.ts_code).split(".")[0] + "_daily_freq_30"
					new_table_class = get_stock_daily_freq(table_name=table_name)
					entry = new_table_class(
						trade_date=trade_date,
						time=row['trade_time'],
						open=row['open'],
						close=row['close'],
						high=row['high'],
						low=row['low'],
						vol=row['vol'],
						amount=row['amount'],
					)
					print(entry.to_dict())
					self.db_main.insert_entry(entry)


if __name__ == '__main__':
	current_date = datetime.now()
	if if_run(current_date):
		trade_date = current_date.strftime('%Y-%m-%d')
		print(trade_date)
		main = StockDailyFreqMain()
		# main.init_data()
		main.task_data(trade_date)

# df = xcsc_pro.stk_mins(ts_code='600095.SH', start_time='2022-08-24 09:30:00', end_time='2023-08-24 15:00:00',
# 					   freq='30min')
# print(df)