from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import json
import pandas as pd
from . import utils


class EntrustOrder(object):
    def __init__(self, security, action, price, size):
        self.security  = security
        self.action    = action
        self.price     = price
        self.size      = size


def set_log_dir(log_dir):
    try:
        import jrpc
        if log_dir:
            jrpc.set_log_dir(log_dir)
        else:
            jrpc.set_log_dir("")
    except Exception as e:
        print("Exception", e)


class TradeApi(object):
    def __init__(self, addr, use_jrpc=False, prod_type="jzts"):
        """
            use_jrpc:
                True     -- Use jrcp_client of C version, for jzts only
                False    -- Use pure python version
            prod_type:
                "jaqs"   -- jrpc_msgpack_wth_snappy
                "jzts"   -- jrpc_msgpack
        """

        self._remote = None
        if prod_type == "jzts":
            try:
                if use_jrpc:
                    import jrpc
                    self._remote = jrpc.JsonRpcClient()
                else:
                    from . import jrpc_py
                    self._remote = jrpc_py.JRpcClient(data_format="msgpack")
            except Exception as e:
                print("Exception", e)

            if not self._remote:
                from . import jrpc_py
                self._remote = jrpc_py.JRpcClient(data_format="msgpack")

        else:
            from . import jrpc_py
            self._remote = jrpc_py.JRpcClient(data_format="msgpack_snappy")

        self._remote.on_rpc_callback = self._on_rpc_callback
        self._remote.on_disconnected = self._on_disconnected
        self._remote.on_connected    = self._on_connected
        self._remote.connect(addr)

        self._ordstatus_callback = None
        self._taskstatus_callback = None
        self._internal_order_callback = None
        self._trade_callback = None
        self._on_connection_callback = None
        self._connected   = False
        self._username    = ""
        self._password    = ""
        self._strategy_id = 0
        self._strategy_selected = False
        self._data_format = "default"

    def __del__(self):
        self._remote.close()
    
    def _on_rpc_callback(self, method, data):
        #print "_on_rpc_callback:", method, data

        if method == "oms.orderstatus_ind":
            if self._data_format == "obj":
                data = utils.to_obj("Order", data)

            if self._ordstatus_callback:
                self._ordstatus_callback(data)

        elif method == "oms.taskstatus_ind":
            if self._data_format == "obj":
                data = utils.to_obj("TaskStatus", data)

            if self._taskstatus_callback:
                self._taskstatus_callback(data)

        elif method == "oms.trade_ind":
            if self._data_format == "obj":
                data = utils.to_obj("Trade", data)

            if self._trade_callback:
                self._trade_callback(data)

        elif method == "oms.internal_order_ind":
            if self._data_format == "obj":
                data = utils.to_obj("QuoteOrder", data)

            if self._internal_order_callback:
                self._internal_order_callback(data)

    def _on_disconnected(self):
        print("TradeApi: _on_disconnected")
        self._connected = False
        self._strategy_selected = False
        if self._on_connection_callback:
            self._on_connection_callback(False)

    def _on_connected(self):
        print("TradeApi: _on_connected")
        self._connected = True
        self._do_login()
        self._do_use_strategy()
        if self._on_connection_callback:
            self._on_connection_callback(True)

    def _check_session(self):
        if not self._connected:
            return (False, "no connection")

        if self._strategy_selected:
            return (True, "")

        r, msg = self._do_login()
        if not r: return (r, msg)
        if self._strategy_id:
            return self._do_use_strategy()
        else:
            return (r, msg)

    def set_data_format(self, format):
        self._data_format = format

    def set_connection_callback(self, callback):
        self._on_connection_callback = callback

    def set_ordstatus_callback(self, callback):
        self._ordstatus_callback = callback

    def set_trade_callback(self, callback):
        self._trade_callback = callback

    def set_task_callback(self, callback):
        self._taskstatus_callback = callback

    def set_quoteorder_callback(self, callback):
        self._internal_order_callback = callback

    def _get_format(self, format, default_format):
        if format:
            return format
        elif self._data_format != "default":
            return self._data_format
        else:
            return default_format

    def login(self, username, password, format=""):
        self._username = username
        self._password = password
        return self._do_login(format=format)

    def _do_login(self, format=""):
        # Shouldn't check connected flag here. ZMQ is a mesageq queue!
        # if !self._connected :
        #    return (False, "-1,no connection")

        if self._username and self._password:
            rpc_params = {"username": self._username,
                          "password": self._password}

            cr = self._remote.call("auth.login", rpc_params)
            f = self._get_format(format, "")
            if f != "obj" and f != "":
                f = ""
            return utils.extract_result(cr, data_format=f, class_name="UserInfo")
        else:
            return (False, "-1,empty username or password")

    def logout(self):
        rpc_params = {}
        
        cr = self._remote.call("auth.logout", rpc_params)
        return utils.extract_result(cr)

    def close(self):
        self._remote.close()

    def use_strategy(self, strategy_id):
        if strategy_id:
            self._strategy_id = strategy_id
            return self._do_use_strategy()
        else:
            # Query
            rpc_params = {"account_id": 0}

            cr = self._remote.call("auth.use_strategy", rpc_params)
            r, msg = utils.extract_result(cr)
            self._strategy_selected = r

            return (r, msg)

    def _do_use_strategy(self):
        if self._strategy_id:
            rpc_params = {"account_id": self._strategy_id}

            cr = self._remote.call("auth.use_strategy", rpc_params)
            r, msg = utils.extract_result(cr)
            self._strategy_selected = r

            return (r, msg)
        else:
            return (False, "-1,no strategy_id was specified")

    def confirm_internal_order(self, task_id, confirmed):
        """
        return (result, message)
        if result is None, message contains error information
        """

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {"task_id"  : task_id,
                      "confirmed": confirmed}

        cr = self._remote.call("oms.confirm_internal_order", rpc_params)
        return utils.extract_result(cr)

    def order(self, security, price, size, algo="", algo_param={}, userdata=""):
        """
        return (result, message)
        if result is None, message contains error information
        """

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {"security"    : security,
                      "price"       : price,
                      "size"        : int(size),
                      "algo"        : algo,
                      "algo_param"  : json.dumps(algo_param),
                      "user"        : self._username,
                      "userdata"    : userdata}

        cr = self._remote.call("oms.order", rpc_params)
        return utils.extract_result(cr)

    def place_order(self, security, action, price, size, algo="", algo_param={}, userdata=""):
        """
        return (result, message)
        if result is None, message contains error information
        """

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = { "security"    : security,
                       "action"      : action,
                       "price"       : price,
                       "size"        : int(size),
                       "algo"        : algo,
                       "algo_param"  : json.dumps(algo_param),
                       "user"        : self._username,
                       "userdata"    : userdata}

        cr = self._remote.call("oms.place_order", rpc_params)
        return utils.extract_result(cr)

    def batch_order(self, orders, algo="", algo_param={}, userdata=""):
        """
        orders format:
            [ {"security": "000001.SZ", "action": "Buy", "price": 10.0, "size" : 100}, ... ]
        return (result, message)
        if result is None, message contains error information
        """

        if not orders or not isinstance(orders, (list, tuple)):
            return (None, "empty order")

        if isinstance(orders[0], EntrustOrder):
            tmp = []
            for o in orders:
                tmp.append({"security": o.security,
                            "price"   : o.price,
                            "size"    : int(o.size)})

            orders = tmp

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {"orders"     : orders,
                      "algo"       : algo,
                      "algo_param" : json.dumps(algo_param),
                      "user"       : self._username,
                      "userdata"   : userdata}

        cr = self._remote.call("oms.batch_order", rpc_params)
        return utils.extract_result(cr)

    def place_batch_order(self, orders, algo="", algo_param={}, userdata=""):
        """
        orders format:
            [ {"security": "000001.SZ", "action": "Buy", "price": 10.0, "size" : 100}, ... ]
        return (result, message)
        if result is None, message contains error information
        """

        if not orders or not isinstance(orders, (list, tuple)):
            return (None, "empty order")

        if isinstance(orders[0], EntrustOrder):
            tmp = []
            for o in orders:
                tmp.append({"security": o.security,
                            "action"  : o.action,
                            "price"   : o.price,
                            "size"    : int(o.size)})

            orders = tmp

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {"orders"     : orders,
                      "algo"       : algo,
                      "algo_param" : json.dumps(algo_param),
                      "user"       : self._username,
                      "userdata"   : userdata}

        cr = self._remote.call("oms.place_batch_order", rpc_params)
        return utils.extract_result(cr)

    def cancel_order(self, task_id):
        """
        return (result, message)
        if result is None, message contains error information
        """

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {"task_id": task_id}

        cr = self._remote.call("oms.cancel_order", rpc_params)
        return utils.extract_result(cr)

    def query_account(self, format=""):
        """
            return pd.dataframe
        """
        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {}

        data_format = self._get_format(format, "pandas")
        if data_format == "pandas":
            rpc_params["format"] = "columnset"

        cr = self._remote.call("oms.query_account", rpc_params)

        return utils.extract_result(cr, data_format=data_format, class_name="Account")

    def query_position(self, mode="all", securities="", format=""):
        """
            securities: seperate by ","
            return pd.dataframe
        """

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {"mode"       : mode,
                      "security"   : securities}

        data_format = self._get_format(format, "pandas")
        if data_format == "pandas":
            rpc_params["format"] = "columnset"

        cr = self._remote.call("oms.query_position", rpc_params)

        return utils.extract_result(cr, data_format=data_format, class_name="Position")

    def query_net_position(self, mode="all", securities="", format=""):
        """
            securities: seperate by ","
            return pd.dataframe
        """

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {"mode"       : mode,
                      "security"   : securities}

        data_format = self._get_format(format, "pandas")
        if data_format == "pandas":
            rpc_params["format"] = "columnset"

        cr = self._remote.call("oms.query_net_position", rpc_params)

        return utils.extract_result(cr, data_format=data_format, class_name="NetPosition")

    def query_repo_contract(self, format=""):
        """
            securities: seperate by ","
            return pd.dataframe
        """

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {}

        cr = self._remote.call("oms.query_repo_contract", rpc_params)

        return utils.extract_result(cr, data_format=self._get_format(format, "pandas"), class_name="RepoContract")

    def query_task(self, task_id=-1, format=""):
        """
            task_id: -1 -- all
            return pd.dataframe
        """

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {"task_id": task_id}

        data_format = self._get_format(format, "pandas")
        if data_format == "pandas":
            rpc_params["format"] = "columnset"

        cr = self._remote.call("oms.query_task", rpc_params)

        return utils.extract_result(cr, data_format=data_format, class_name="Task")

    def query_order(self, task_id=-1, format=""):
        """
            task_id: -1 -- all
            return pd.dataframe
        """

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {"task_id": task_id}

        data_format = self._get_format(format, "pandas")
        if data_format == "pandas":
            rpc_params["format"] = "columnset"

        cr = self._remote.call("oms.query_order", rpc_params)

        return utils.extract_result(cr, data_format=data_format, class_name="Order")

    def query_trade(self, task_id=-1, format=""):
        """
            task_id: -1 -- all
            return pd.dataframe
        """

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {"task_id": task_id}

        data_format = self._get_format(format, "pandas")
        if data_format == "pandas":
            rpc_params["format"] = "columnset"

        cr = self._remote.call("oms.query_trade", rpc_params)

        return utils.extract_result(cr, data_format=data_format, class_name="Trade")

    def query_portfolio(self, format=""):
        """
            return pd.dataframe
        """

        r, msg = self._check_session()
        if not r: return (None, msg)

        rpc_params = {}

        data_format = self._get_format(format, "pandas")
        if data_format == "pandas":
            rpc_params["format"] = "columnset"

        cr = self._remote.call("pms.query_portfolio", rpc_params)

        return utils.extract_result(cr, index_column="security", data_format=data_format, class_name="NetPosition")

    def goal_portfolio(self, positions, algo="", algo_param={}, userdata=""):
        """
        positions format:
            [ {"security": "000001.SZ", "ref_price": 10.0, "size" : 100}, ...]
        return (result, message)
        if result is None, message contains error information
        """

        r, msg = self._check_session()
        if not r: return (False, msg)

        if type(positions) is pd.core.frame.DataFrame :
            tmp = []
            for i in range(0, len(positions)):
                tmp.append ({'security': positions.index[i],
                             'ref_price': float(positions['ref_price'][i]),
                              "size" : int(positions['size'][i])})
            positions = tmp

        rpc_params = {"positions"  : positions,
                      "algo"       : algo,
                      "algo_param" : json.dumps(algo_param),
                      "user"       : self._username,
                      "userdata"   : userdata}

        cr = self._remote.call("pms.goal_portfolio", rpc_params)
        return utils.extract_result(cr)

    def basket_order(self, orders, algo="", algo_param={}, userdata=""):
        """
        orders format:
            [ {"security": "000001.SZ", "ref_price": 10.0, "inc_size" : 100}, ...]
        return (result, message)
        if result is None, message contains error information
        """

        r, msg = self._check_session()
        if not r: return (False, msg)

        if type(orders) is pd.core.frame.DataFrame :
            tmp = []
            for i in range(0, len(orders)):
                tmp.append ({'security': orders.index[i],
                             'ref_price': float(orders['ref_price'][i]),
                             'inc_size' : int(orders['inc_size'][i])})
            orders = tmp

        rpc_params = {"orders": orders,
                      "algo": algo,
                      "algo_param": json.dumps(algo_param),
                      "user": self._username,
                      "userdata": userdata}

        cr = self._remote.call("pms.basket_order", rpc_params)
        return utils.extract_result(cr)

    def stop_portfolio(self):
        """
        return (result, message)
        if result is None, message contains error information
        """

        r, msg = self._check_session()
        if not r: return (False, msg)

        rpc_params = {}
    
        cr = self._remote.call("pms.stop_portfolio", rpc_params)
        return utils.extract_result(cr)

    def query_universe(self, format=""):

        r, msg  = self._check_session()
        if not r: return (None, msg)

        rpc_params = {}
        data_format = self._get_format(format, "pandas")
        if data_format == "pandas":
            rpc_params["format"] = "columnset"
        
        cr = self._remote.call("oms.query_universe", rpc_params)
        
        return utils.extract_result(cr, data_format=data_format, class_name="UniverseItem")
    
    def set_heartbeat(self, interval, timeout):
        self._remote.set_hearbeat_options(interval, timeout)
        print("heartbeat_interval =", self._remote._heartbeat_interval, ", heartbeat_timeout =",
              self._remote._heartbeat_timeout)
