# -*- coding: utf-8 -*-
"""
Created on Fri Aug  5 15:29:12 2022

@author: MINGL
"""
import pandas as pd
import pyautogui as pyautogui
import pywinauto
from pywinauto import clipboard
from pywinauto import keyboard
from skimage import io
import time
import io as mio
import pytesseract
import warnings
import datetime
from tools import *


from connect_wifi import connect_wifi
from iptest import get_proxy_ip,set_global_proxy,get_ip_data

warnings.filterwarnings('ignore')
scripts_path = os.path.dirname(os.path.realpath(__file__))
root_path = scripts_path[:scripts_path.find(project_name)+len(project_name)]

path = r'C:\trade_software1\trade_software\zd_zsone\TdxW.exe'


class ZSClientTrader:
    """
    基于通达信程序国信金太阳专业版自动交易程序
    """

    def __init__(self, account_name, exe_path=path):
        self.path = exe_path
        self.account_name = account_name
        self.log = Logger(f'{root_path}/logs',self.account_name)
        self.As_Of_Date = datetime.datetime.today().strftime('%Y%m%d')
        # 用户券商信息
        sql = f"select * from ainvest_usercount where username='{account_name}'"
        df_count = download_data_from_db(sql, 'ai_strategy_update_iddb')
        self.securities_name = str(df_count['securities_username'].tolist()[0])
        self.securities_password = str(df_count['securities_password'].tolist()[0])
        self.communication_password = str(df_count['communication_password'].tolist()[0])
        self.trade_control_id_group = {
            's1': 0x3FD, 's2': 0x3FE, 's3': 0x3FF, 's4': 0x409, 's5': 0x408,
            'b1': 0x3FA, 'b2': 0x401, 'b3': 0x402, 'b4': 0x40B, 'b5': 0x40C,
        }
    def exit(self):
        '''
        结束进程

        Returns
        -------
        None.

        '''
        self._app.kill()
    def _close_prompt_windows(self):
        time.sleep(0.1)
        for window in self._app.windows(class_name="#32770"):
            if window.window_text() != '网上股票交易系统*':
                window.close()
        time.sleep(0.1)
    def login(self):
        '''
        启动软件并登录

        Returns
        -------
        None.

        '''
        #软件启动判断，若有进行中进程，先关闭再重启
        try:
            self._app = pywinauto.Application().connect(
                path=self.path, timeout=1
            )
            self.exit()
            self.log.warn('There are running programs, and there are processes in retreat')
        except Exception:
            pass

        #启动软件进程，输入账号密码
        self.log.info('Start login')
        self._app = pywinauto.Application().start(self.path)
        while True:
            try:
                self._app.window(title_re='招商证券.*').wait("ready")
                break
            except RuntimeError:
                pass
        time.sleep(1)
        self._app.window(title_re='招商证券.*').window(control_id = 0xE9,class_name = 'SafeEdit').click()
        pyautogui.typewrite(self.securities_name)
        time.sleep(1)
        self._app.window(title_re='招商证券.*').window(control_id = 0xEA,class_name = 'SafeEdit').click()
        pyautogui.typewrite(self.securities_password)
        # 登录
        pywinauto.keyboard.send_keys("{ENTER}")

        # while True:
        #     try:
        #         # self._app['国信证券消息中心'].wait("ready")
        #         self._app['招商证券智远一户通V6.72 - [版面-自选股股]'].wait("ready")
        #         break
        #     except RuntimeError:
        #         pass
        time.sleep(12)
        tryrun2(self._app['招商证券智远一户通V6.72 - [版面-自选股股]'].wait("ready"))
        # time.sleep(2)
        #关闭弹窗
        # self._close_prompt_windows()

        for i in range(3):
            self._turn_off_advertising()
        # self._close_prompt_windows()
        #设置主窗口及菜单窗口
        self.main_wnd = self._app.TdxW_MainFrame_Class
        self.left_wnd = self._app.TdxW_MainFrame_Class.window(class_name='SysTreeView32', control_id=0xE900)

    def _turn_off_advertising(self):
        '''
        关闭弹窗

        Returns
        -------
        None.

        '''
        try:
            self._app.top_window().set_focus()
            time.sleep(0.5)
            pywinauto.keyboard.send_keys("{ESC}")
            pywinauto.keyboard.send_keys("{ENTER}")
            pywinauto.keyboard.send_keys("{ESC}")
        except Exception as e:
            pass
    def buy(self, stock_no, price, num):
        '''
        买入股票

        Parameters
        ----------
        stock_no : TYPE
            DESCRIPTION.
        price : TYPE
            DESCRIPTION.
        num : TYPE
            DESCRIPTION.

        Returns
        -------
        None.

        '''
        time.sleep(0.5)
        path = ['买入']
        # 点击买入
        self.left_wnd.get_item(path, exact=False).click_input(button=u'left', double=False, wheel_dist=0,
                                                              pressed=u'')
        time.sleep(1)
        self.__trade(stock_no, price, num)

    def sell(self, stock_no, price, num):
        '''
        卖出股票

        Parameters
        ----------
        stock_no : TYPE
            DESCRIPTION.
        price : TYPE
            DESCRIPTION.
        num : TYPE
            DESCRIPTION.

        Returns
        -------
        None.

        '''
        time.sleep(0.5)
        path = ['卖出']
        # 点击卖出
        self.left_wnd.get_item(path, exact=False).click_input(button=u'left', double=False, wheel_dist=0,
                                                              pressed=u'')
        time.sleep(1)
        self.__trade(stock_no, price, num)

    def purchase_stocks(self):

        self.left_wnd.get_item([5,2]).click_input()
        time.sleep(1)
        self.main_wnd.window(control_id=0x9876,class_name='Button').click()

        one_text = str(self._app.window(control_id=0x0, class_name="#32770",found_index=0).children()[4].context_help_id).split("'")[1].split("'")[0]
        self.log.info(one_text.replace('/n', ' '))

        while True:
            try:
                self._app.window(control_id=0x0, class_name="#32770", found_index=0).window(control_id=0x1B67,class_name="Button").click_input()

                self.log.info(str(self._app.window(title='提示', class_name="#32770", found_index=0).children()[4].context_help_id).split("'")[1].split("'")[0])

            except:
                break

    def cancel_entrust(self):
        '''
        全撤

        Returns
        -------
        None.

        '''
        # # 进入撤单页面
        time.sleep(0.5)
        self.left_wnd.get_item([4]).click_input()        # 点击撤单
        time.sleep(1)
        # 刷新
        try:
            self.main_wnd.window(control_id=0x474, class_name='Button').click_input()
        except:
            self.main_wnd.window(control_id=0x474, class_name='Button').click()

        time.sleep(1)
        try:
            # 全选中
            self.main_wnd.window(control_id=0xE, class_name='Button').click()
        except:
            self.main_wnd.window(control_id=0xE, class_name='Button').click_input()

        time.sleep(1)
        # exit()
        # 全撤
        try:
            self.main_wnd.window(control_id=0x470, class_name='Button').click_input()
        except:
            self.main_wnd.window(control_id=0x470, class_name='Button').click()

        time.sleep(1)

        # 确认撤单
        pywinauto.keyboard.send_keys("{ENTER}")
        time.sleep(5)
        # 确认撤单2
        pywinauto.keyboard.send_keys("{ENTER}")
        time.sleep(5)
        # 确认撤单2
        pywinauto.keyboard.send_keys("{ESC}")
        time.sleep(0.5)
        pywinauto.keyboard.send_keys("{ENTER}")

    def _delete_first_lines(self,filename,pp):
        fin = open(filename, 'r')
        a = fin.readlines()
        fout = open(filename, 'w')
        if pp==0:
            b = ''.join(a[:3])
        else:
            b = ''.join(a[3:])
        fout.write(b)

    def get_balance(self):
        '''
        获取资金情况

        Returns
        -------
        record : TYPE
            DESCRIPTION.

        '''
        self.left_wnd.get_item([3, 0]).click_input()
        time.sleep(1)
        self._copy_data([3, 0])
        filenme = f'C:\\Users\\Administrator\\Documents\\{self.As_Of_Date} 资金股份查询.xls'
        self.df_balance = pd.read_csv(filenme, encoding='gbk', sep='\t', engine='python', error_bad_lines=False)
        # print(self.df_balance)
        self.df_balance['="可用"'] = self.df_balance['="可用"'].map(lambda x: x.strip('=').strip('"'))
        self.df_balance['="参考市值"'] = self.df_balance['="参考市值"'].map(lambda x: x.strip('=').strip('"'))
        self.df_balance['="资产"'] = self.df_balance['="资产"'].map(lambda x: x.strip('=').strip('"'))

        result = {}
        result['Cash'] = list(self.df_balance['="可用"'].values)[0]
        result['Position_amount'] = list(self.df_balance['="参考市值"'].values)[0]
        result['Total_account'] = list(self.df_balance['="资产"'].values)[0]
        result = pd.DataFrame(result, index=[0])
        result['Account_Name'] = self.account_name
        result['As_Of_Date'] = datetime.date.today()
        return result

    def get_position(self):
        '''
        获取持仓情况

        Returns
        -------
        df_position : TYPE
            DESCRIPTION.

        '''
        self.left_wnd.get_item([3, 0]).click_input()
        # 复制表格数据
        time.sleep(1)
        self._copy_data([3, 0])


        filenme = f'C:\\Users\\Administrator\\Documents\\{str(datetime.date.today()).replace("-","")} 资金股份查询.xls'
        self._delete_first_lines(filenme, 1)
        # 读文件
        df_position = pd.read_csv(filenme, encoding='gbk', sep='\t', engine='python', error_bad_lines=False)
        os.remove(filenme)
        if len(df_position) == 0:
            return pd.DataFrame()
        else:
            df_position.columns = [i.strip('=').strip('"') for i in list(df_position.columns)]
            df_position = df_position.applymap(lambda x: x.strip('=').strip('"') if isinstance(x, str) else x)
            df_position['证券代码'] = df_position['证券代码'].apply(lambda x: tranTicker(x))
            df_position = df_position[
                ['证券代码', '证券名称', '证券数量', '可卖数量',  '成本价', '当前价', '最新市值',
                 '持仓盈亏', '持仓盈亏比例(%)']]

            df_position.rename(columns={"证券代码": "Ticker",
                                        "证券名称": "Ticker_name",
                                        "证券数量": "Number_transactions",
                                        "可卖数量": "cash",
                                        "成本价": "cost_price",
                                        "当前价": "market_price",
                                        "最新市值": "market_value",
                                        "持仓盈亏": "profit_and_loss",
                                        "持仓盈亏比例(%)": "profit_loss_ratio",
                                        }, inplace=True)
            df_position['frozen_quantity'] = 0
            df_position['As_Of_Date'] = datetime.date.today()
            df_position['Account_Name'] = self.account_name

            return df_position



    def get_order(self):
        '''
        获取委托情况

        Returns
        -------
        TYPE
            DESCRIPTION.

        '''
        self.left_wnd.get_item([6, 1]).click_input()
        # 复制表格数据
        time.sleep(1)
        self._copy_data([6, 1])
        try:
            filenme = f'C:\\Users\\Administrator\\Documents\\{str(datetime.date.today()).replace("-", "")} 当日委托查询.xls'
            self._delete_first_lines(filenme, 0)
            # 读文件
            df_order = pd.read_csv(filenme, encoding='gbk', sep='\t', engine='python', error_bad_lines=False)
            os.remove(filenme)
        except:
            self.log.info('--------------------------------------无当日委托--------------------------------------')
            return pd.DataFrame()
        df_order.columns = [i.strip('=').strip('"') for i in list(df_order.columns)]
        df_order = df_order.applymap(lambda x: x.strip('=').strip('"') if isinstance(x, str) else x)
        df_order['证券代码'] = df_order['证券代码'].apply(lambda x: tranTicker(x))
        return df_order

    def get_today_trades(self):
        '''
        获取当日成交

        Returns
        -------
        TYPE
            DESCRIPTION.

        '''
        self.left_wnd.get_item([3, 2]).click_input()
        # 复制表格数据
        time.sleep(1)
        self._copy_data([3, 2])
        try:
            filenme = f'C:\\Users\\Administrator\\Documents\\{str(datetime.date.today()).replace("-", "")} 当日成交查询.xls'
            self._delete_first_lines(filenme, 0)
            # 读文件
            df = pd.read_csv(filenme, encoding='gbk', sep='\t', engine='python', error_bad_lines=False)
            os.remove(filenme)
            df.columns = [i.strip('=').strip('"') for i in list(df.columns)]
            df = df.applymap(lambda x: x.strip('=').strip('"') if isinstance(x, str) else x)
            df['证券代码'] = df['证券代码'].apply(lambda x: tranTicker(x))

            df = df[['成交时间', '证券代码', '证券名称', '买卖标志', '成交数量', '成交价格']]
            df['Account_Name'] = self.account_name
            df['As_Of_Date'] = datetime.date.today()
            df.rename(columns={'成交时间': 'Trade_time', '证券代码': 'Ticker', '证券名称': 'Ticker_name',
                               '买卖标志': 'Operate',
                               '成交数量': 'Number_transactions',
                               '成交价格': 'Average_price'}, inplace=True)
            df['Operate'] = df['Operate'].apply(lambda x: '卖' if '卖' in x else '买')

            df['Ticker'] = df['Ticker'].apply(lambda x: tranTicker(str(x).replace('="', '').replace('"', '')))

            return df
        except:
            self.log.info('--------------------------------------无当日成交--------------------------------------')
            return pd.DataFrame()


    def __trade(self, stock_no, price, amount):
        '''
        交易输入函数

        Parameters
        ----------
        stock_no : TYPE
            DESCRIPTION.
        price : TYPE
            DESCRIPTION.
        amount : TYPE
            DESCRIPTION.

        Returns
        -------
        None.

        '''
        time.sleep(1)
        # 设置股票代码
        pyautogui.typewrite(str(stock_no))
        time.sleep(0.5)

        if price in ['b1', 'b2', 'b3', 'b4', 'b5', 's1', 's2', 's3', 's4', 's5']:
            price = float(
                self.main_wnd.window(control_id=self.trade_control_id_group[price], class_name="Static").window_text())

        # 设置价格
        # pyautogui.typewrite(str(price), interval=0.1)
        self.main_wnd.window(control_id=0x2EE6, class_name="Edit").set_text(str(price))
        # self.main_wnd.Edit4.set_text(str(price))
        time.sleep(1)

        # 设置股数目
        self.main_wnd.window(control_id=0x2EE7, class_name="Edit").set_text(str(amount))
        time.sleep(1)
        # exit()

        # 点击卖出or买入
        try:
            self.main_wnd.window(control_id=0x7DA, class_name="Button").click()
        except:
            self.main_wnd.window(control_id=0x7DA, class_name="Button").click_input()

        time.sleep(1)

        # # 摘取合同信息
        # # result_1 = self._app['买入交易确认'].child_window(class_name='Static',control_id=0x1B65).window_text()
        # result_1 = self._app['买入交易确认'].window_text()
        # self.__parse_contract(result_1)
        # time.sleep(0.2)

        # 确认买入
        # pywinauto.keyboard.send_keys("{ENTER}")
        self._app.window(class_name='#32770', control_id=0x0, found_index=0).window(control_id=0x1B67,
                                                                                    class_name='Button').click()
        time.sleep(1)

        # # 摘取合同信息2
        # # result_2 = self._app['提示'].child_window(control_id=0x1B65, class_name='Static').window_text()
        # result_2 = self._app['提示'].Static1.window_text()
        # print(result_2)

        # 确认合同
        # pywinauto.keyboard.send_keys("{ENTER}")
        self._app.window(class_name='#32770', control_id=0x0, found_index=0).window(control_id=0x1B67,
                                                                                    class_name='Button').click()
        time.sleep(0.5)


    def _copy_data(self, path):
        self.left_wnd.get_item(path, exact=False).click_input(button=u'left', double=False, wheel_dist=0, pressed=u'')
        time.sleep(0.5)
        try:
            self.main_wnd.window(control_id=0x47F, class_name='Button').click()
            time.sleep(0.5)

            self._app['输出'].window(control_id=0xE9, class_name='Button').click()
        except Exception:
            print('无输出的数据！')
        # 确认
        pywinauto.keyboard.send_keys("{ENTER}")

    def _delete_first_lines(self, filename, pp):

        fin = open(filename, 'r')
        a = fin.readlines()
        fout = open(filename, 'w')
        if pp == 0:
            b = ''.join(a)
        else:
            b = ''.join(a[3:])
        fout.write(b)

if __name__ == '__main__':

    account_name = str('13240313541')
    a = ZSClientTrader(account_name)
    a.login()

    # 买入
    # a.buy('600000', '7','100')

    # 卖出
    # a.sell('600835', '12.45','900')
    # print(a.get_order())

    # a.purchase_stocks()

    # # 获取当日实际持仓
    # print(a.get_position())
    #
    # # 获取资金情况
    # print(a.get_balance())
    #
    # # 获取当日成交
    # print(a.get_today_trades())
