# auto generated by update_py.py

import csv
import datetime
import json
import os
import queue
import requests
import sys
import time

from concurrent.futures import ThreadPoolExecutor
from influxdb import InfluxDBClient
from influxdb.exceptions import InfluxDBClientError, InfluxDBServerError

import tlclient.trader.message_common as message_common
import tlclient.trader.message_market as message_market
import tlclient.trader.message_trade as message_trade

from tlclient.linker.constant import FistType
from tlclient.linker.fist import Fist
from tlclient.linker.utility import get_today_date
from tlclient.trader.constant import MsgType


class RecorderClient:

    def __init__(self, logger):
        self.logger = logger
        self._to_stop = False

    def writes(self, msg_obj, f):
        pass

    def stop(self):
        self._to_stop = True
        self.close_client()

    def close_client(self):
        pass

    def is_running(self):
        return not self._to_stop


class CsvFileClient(RecorderClient):

    def __init__(self, path, logger):
        RecorderClient.__init__(self, logger)
        self.base_path = path
        self.file_clients = dict()

    def _write_header(self, msg_class_name, msg_obj_dict):
        self.file_clients[msg_class_name]['fwrite'].writeheader()

    def _write_content(self, msg_class_name, msg_obj_dict):
        # suppose dict is in order (python >= 3.6)
        self.file_clients[msg_class_name]['fwrite'].writerow(msg_obj_dict)

    def _add_file_client(self, msg_class_name, msg_obj_dict, today_date):
        file_path = os.path.join(
            self.base_path, msg_class_name + '_' + today_date + '.csv')
        file_exists = os.path.exists(file_path)
        self.file_clients[msg_class_name] = {
            'date': today_date,
            'file_client': open(file_path, 'a', buffering=1),
            'fwrite': None
        }
        fieldnames = msg_obj_dict.keys()
        self.file_clients[msg_class_name]['fwrite'] = csv.DictWriter(
            self.file_clients[msg_class_name]['file_client'], fieldnames=fieldnames)
        if not file_exists:
            self._write_header(msg_class_name, msg_obj_dict)

    def writes(self, msg_obj, f):
        msg_class_name = msg_obj.__class__.__name__
        if 'to_dict' in msg_obj.__class__.__dict__:
            msg_obj_dict = msg_obj.to_dict()
            self.logger.debug('[write_csv] (msg_obj){}'.format(msg_obj))

            today_date = get_today_date()
            if msg_class_name not in self.file_clients:
                self._add_file_client(msg_class_name, msg_obj_dict, today_date)
            elif self.file_clients[msg_class_name]['date'] != today_date:
                self.file_clients[msg_class_name]['file_client'].close()
                self._add_file_client(msg_class_name, msg_obj_dict, today_date)

            self._write_content(msg_class_name, msg_obj_dict)

    def close_client(self):
        self.logger.debug("[csv_client] closed (base_path)'{}'".format(self.base_path))
        for msg_class_name in self.file_clients.keys():
            self.file_clients[msg_class_name]['file_client'].close()
        self.file_clients.clear()


class KdbClient(RecorderClient):

    def __init__(self, host, port, logger):
        RecorderClient.__init__(self, logger)
        self.addr = "{}:{}".format(host, port)

    def write_kdb(self, table, msg_obj, f, cmd='.q.upsert'):
        data_info = msg_obj.to_kdb()
        q_str = self.dump_frame_info_to_kdb(data_info, f)
        self.logger.debug('[write_kdb] (table){} (q_str){}'.format(table, q_str))
        url = "http://{}/?{}[{};{}]".format(self.addr, cmd, table, q_str)
        requests.get(url)

    def dump_frame_info_to_kdb(self, data_info, f):
        frame_info = [f.get_source(), f.get_req_id(), f.get_err_id(), f.get_msg_type()]
        k_str = self.stringlize_q(data_info + frame_info)
        return k_str

    def writes(self, msg_obj, f):
        msg_type = f.get_msg_type()
        if msg_type == MsgType.MKT_SNAP:
            self.write_kdb("`mktSnap", msg_obj, f, "upsert_snap")
        elif msg_type == MsgType.MKT_TRADE:
            self.write_kdb("`mktTrade", msg_obj, f)
        elif msg_type == MsgType.MKT_BAR:
            self.write_kdb("`mktBar", msg_obj, f, "upsert_bar")
        elif msg_type == MsgType.RSP_POSITION:
            self.write_kdb("`positions", msg_obj, f)

    def stringlize_q(self, data):
        if isinstance(data, list):
            q_format = "({})".format(";".join([self.stringlize_q(i) for i in data]))
            return q_format
        elif isinstance(data, int):
            return "{}i".format(data)
        elif isinstance(data, float):
            return "{}f".format(data)
        elif isinstance(data, datetime.date):
            return data.strftime("%Y.%m.%d")
        elif isinstance(data, datetime.time):
            return data.strftime('%H:%M:%S.%f')[:-3]
        elif isinstance(data, str) and ";" not in data:
            return '`$("{}")'.format(data)
        return str(data)

    def close_client(self):
        ## todo:
        ## close the kdb handler
        pass


class InfluxClient(RecorderClient):

    def __init__(self, host, port, logger, db_name, udp_port):
        RecorderClient.__init__(self, logger)

        self.client = InfluxDBClient(host=host, port=port, udp_port=udp_port, use_udp=(udp_port is not None), database=db_name, retries=100)
        self.client.create_database(db_name)

        self.thread_pool = ThreadPoolExecutor(20, thread_name_prefix='influx_thread_pool')
        self.batch_size = 20 if udp_port is not None else 5000
        self.write_queue = queue.Queue()

        self.thread_pool.submit(self._write_queue_checker)

    def _proc_write_thread(self, res):
        try:
            start = time.time()
            self.logger.debug('[influxdb] start write_points. (res){}'.format(len(res)))
            self.client.write_points(res, time_precision='u', batch_size=self.batch_size)
            writes_time = time.time() - start
            if writes_time < 1:
                self.logger.debug('[influxdb] end write_points. (res){} (time){}'.format(len(res), writes_time))
            elif writes_time < 2:
                self.logger.info('[influxdb] end write_points. (res){} (time){}'.format(len(res), writes_time))
            else:
                self.logger.warning('[influxdb] end write_points. (res){} (time){}'.format(len(res), writes_time))

        except InfluxDBClientError as e:
            err_msg = json.loads(e.content)
            if err_msg.get('error', '').startswith('partial write: points beyond retention policy dropped'):
                self.logger.warning('[influxdb] write_points client error. (err_msg){}'.format(err_msg['error']))
            else:
                self.logger.error('[influxdb] write_points client error, will stop influx client. (err_msg){}'.format(err_msg))
                self.stop()

        except InfluxDBServerError:
            self.logger.error('[influxdb] write_points error, will stop influx client. (err_msg){}'.format(sys.exc_info()))
            self.stop()

    def _write_queue_checker(self):
        while self.is_running() or not self.write_queue.empty():
            res = []
            i = 0
            while i < 5000:
                try:
                    item = self.write_queue.get(timeout=0.1)
                    res.append(item)
                    i += 1
                except queue.Empty:
                    break
            if len(res) > 0:
                self.thread_pool.submit(self._proc_write_thread, res)

    def dump_frame_info_to_flux(self, flux_info, f):
        if isinstance(flux_info, list):
            return [self.dump_frame_info_to_flux(flux_item, f) for flux_item in flux_info]
        else:
            flux_info['tags'].update({
                'frame_source': f.get_source(),
                'frame_req_id': f.get_req_id(),
                'frame_err_id': f.get_err_id(),
                'frame_msg_type': f.get_msg_type()
            })
            flux_info['fields'].update({
                'frame_dt': f.get_nano()
            })
            return flux_info

    def writes(self, msg_obj, f):
        if hasattr(msg_obj, 'to_influx'):
            flux_info = msg_obj.to_influx()
            final_info = self.dump_frame_info_to_flux(flux_info, f)
            if not isinstance(final_info, list):
                final_info = [final_info]
            for info in final_info:
                if info.get('measurement'):
                    self.write_queue.put(info)

    def close_client(self):
        self.client.close()


class Recorder(Fist):

    def __init__(self, name, env_name, addr):
        Fist.__init__(self, name, FistType.RECORDER, env_name)
        self.set_master_addr(addr)
        self.create_fist()
        self.clients = []
        self._msg_callbacks = {}

    def sub_router(self, router_name):
        self.reg_sub(router_name)

    ######################
    # internal functions #
    ######################

    def register_msg_callback(self, msg_type, callback):
        # callback should have the args list exactly matched (data: dict, msg_type: int, frame: Frame)
        if msg_type in self._msg_callbacks:
            self.logger.warn('[reg_msg_callback] replacing callback (msg_type){}/{}'.format(msg_type, MsgType.read(msg_type)))
        self._msg_callbacks[msg_type] = callback

    def on_pub_frame(self, f):
        msg_type = f.get_msg_type()
        msg_obj = None

        msg_callback = self._msg_callbacks.get(msg_type)
        if msg_callback is not None:
            data = json.loads(f.get_string())
            msg_callback(data, msg_type, f)

        elif MsgType.is_market_data_type(msg_type):
            if msg_type == MsgType.MKT_SNAP:
                msg_obj = f.get_obj(message_market.MktSnap)

            elif msg_type == MsgType.MKT_SNAP_PLUS:
                msg_obj = f.get_obj(message_market.MktSnapPlus)

            elif msg_type == MsgType.MKT_SNAP_FUT:
                msg_obj = f.get_obj(message_market.MktSnapFut)

            elif msg_type == MsgType.MKT_SNAP_OPT:
                msg_obj = f.get_obj(message_market.MktSnapOpt)

            elif msg_type == MsgType.MKT_BAR_GEN:
                msg_obj = f.get_obj(message_market.MktBarGen)

            elif msg_type == MsgType.MKT_INDEX:
                msg_obj = f.get_obj(message_market.MktIndex)

            elif msg_type == MsgType.MKT_ORDER:
                msg_obj = f.get_obj(message_market.MktOrder)

            elif msg_type == MsgType.MKT_TRADE:
                msg_obj = f.get_obj(message_market.MktTrade)

            elif msg_type == MsgType.MKT_BAR:
                msg_obj = f.get_obj(message_market.MktBar)

            elif msg_type == MsgType.MKT_BAR_GEN:
                msg_obj = f.get_obj(message_market.MktBarGen)

            elif msg_type == MsgType.MKT_VOL:
                msg_obj = f.get_obj(message_market.MktVol)

        elif MsgType.is_trading_data_type(msg_type):
            if msg_type == MsgType.RSP_ORDER_INSERT:
                msg_obj = f.get_obj(message_trade.RspOrderInsert)

            elif msg_type == MsgType.RSP_ORDER_CANCEL:
                msg_obj = f.get_obj(message_trade.RspOrderCancel)

            elif msg_type == MsgType.RTN_ORDER:
                msg_obj = f.get_obj(message_trade.RtnOrder)

            elif msg_type == MsgType.RTN_TRADE:
                msg_obj = f.get_obj(message_trade.RtnTrade)

            elif msg_type == MsgType.RSP_POSITION:
                pos_d = json.loads(f.get_string())
                msg_obj = message_trade.RspPosition(pos_d)

            elif msg_type == MsgType.RSP_ACCOUNT:
                msg_obj = f.get_obj(message_trade.RspAccount)

            elif msg_type == MsgType.RSP_ACTIVE_ORDERS:
                ods_d = json.loads(f.get_string())
                msg_obj = message_trade.RspActiveOrders(ods_d)

            elif msg_type == MsgType.RSP_CANCEL_ACTIVE_ORDERS:
                msg_obj = f.get_obj(message_trade.RspCancelActiveOrders)

        elif MsgType.is_system_status_data_type(msg_type):
            if msg_type == MsgType.GTW_CONNECTION:
                msg_obj = f.get_obj(message_common.GatewayConnectionStatus)

        if msg_obj:
            self.process_msg_obj(msg_obj, f)

    def process_msg_obj(self, msg_obj, frame):
        writes_result = False
        for client in self.clients:
            if client.is_running():
                client.writes(msg_obj, frame)
                writes_result = True

        if not writes_result:
            self.logger.error('[recorder] all clients are not running, will stop recorder.')
            self.stop()

    def add_kdb_client(self, host="localhost", port=5566):
        kdb_client = KdbClient(host, port, self.logger)
        self.clients.append(kdb_client)

    def add_influxdb_client(self, host="localhost", port=8086, db_name='trader', udp_port=None):
        influx_client = InfluxClient(host, port, self.logger, db_name, udp_port)
        self.clients.append(influx_client)

    # csv setting
    def add_csv_client(self, path):
        if not os.path.exists(path):
            os.makedirs(path)
        elif not os.path.isdir(path):
            self.logger.warning("[csv_client] reject by isfile=True. (path){}".format(path))
        csv_file_client_paths = [client.base_path for client in self.clients if isinstance(client, CsvFileClient)]
        normpath = os.path.normpath(path)
        if normpath in csv_file_client_paths:
            self.logger.warning("[csv_client] reject by duplicated. (path){}".format(normpath))
            return

        csv_file_client = CsvFileClient(normpath, self.logger)
        self.clients.append(csv_file_client)

    def on_close(self):
        self.logger.info("[Recorder] stop clients (fist_name){}".format(self.fist_name))
        for client in self.clients:
            client.stop()
        self.clients = []
