from sqlalchemy import and_

from DB.model.StockBasic import StockBasic
from DB.sqlite_db_main import SqliteDbMain, config
from utils.comm import *


class StockInfoMain:

	def __init__(self):
		self.db_main = SqliteDbMain(config.stock_info_db)
		self.session = self.db_main.get_session()

	def insert_stock_basic(self):
		pro = xcsc_pro
		info = ['SZSE', 'SSE']
		# entries = []
		for item in info:
			# SSE: 上交所 SZSE: 深交所
			df = pro.stock_basic(exchange=item)
			for index, row in df.iterrows():
				entry = StockBasic(**row)
				self.db_main.insert_or_update(entry, query_conditions={'ts_code': row['ts_code']})

	def get_stock_basic(self, ts_code=None, symbol=None, restart_id=0):
		sql_cond = and_(
			StockBasic.delist_date.is_('None')
		)
		try:
			if restart_id > 0:
				sql_cond = and_(
					StockBasic.id >= restart_id,
					StockBasic.delist_date.is_('None')
				)
			if ts_code is not None:
				if isinstance(ts_code, str):
					sql_cond = and_(
						StockBasic.ts_code.is_(ts_code),
						StockBasic.delist_date.is_('None')
					)
				elif isinstance(ts_code, list):
					sql_cond = and_(
						StockBasic.ts_code.in_(ts_code),
						StockBasic.delist_date.is_('None')
					)
			elif symbol is not None:
				if isinstance(symbol, str):
					sql_cond = and_(
						StockBasic.symbol.is_(symbol),
						StockBasic.delist_date.is_('None')
					)
				elif isinstance(symbol, list):
					sql_cond = and_(
						StockBasic.ts_code.in_(symbol),
						StockBasic.delist_date.is_('None')
					)
			results = self.session.query(StockBasic).filter(sql_cond).all()
			return results
		finally:
			self.session.close()


# try:
# 	results = self.session.query(StockBasic).filter(StockBasic.delist_date.is_(None)).all()
# 	return results
# finally:
# 	self.session.close()


if __name__ == '__main__':
	print(StockInfoMain().get_stock_basic(restart_id=5))