import json
import threading
from dataclasses import dataclass
from threading import Lock
from typing import Any, Callable, Dict, List, Literal, Optional, Set, Union, get_args

import betterproto
import paho.mqtt.client as mqtt

from settrade.pb.bidofferv3_pb2 import BidOfferV3
from settrade.pb.exchinfodvv3_pb2 import DerivExchangeInfoV3
from settrade.pb.exchinfoeqv3_pb2 import EquityExchangeInfoV3
from settrade.pb.infoV3_pb2 import InfoV3
from settrade.pb.orderdvv3_pb2 import OrderDerivV3
from settrade.pb.ordereqv3_pb2 import OrderEquityV3
from settrade.util import mqtt_to_message

from .context import Context, Option

CallbackType = Literal[
    "on_message",
    "on_disconnect",
    "on_subscribe",
    "on_unsubscribe",
    "on_connect",
    "on_connect_fail",
    "on_log",
    "on_publish",
]

CALLBACK_TYPE_LIST = get_args(CallbackType)


@dataclass
class Callback:
    topic: str
    callback: Callable
    callback_type: CallbackType


class CallBacker:
    is_connected = False
    message_callback_count = 0
    base_topic_count = 0
    lock_callback_pool = Lock()

    def __init__(
        self,
        host: Optional[str],
        path: str,
        token_type: str,
        token: Optional[str],
        client_id: str = "",
    ) -> None:
        """Create mqtt client

        1. initialize client with clean_session=True and transport="websockets"
        clean_session=True mean the broker will remove all information about this client when it disconnects.
        transport="websockets" mean the client will use websocket to connect to broker, please beware that default is tcp.

        2. set tsl by tsl_set() (ssl.PROTOCOL_TLSv1_2) because settrade broker only accept secure connection
        if you got a message like "An existing connection was forcibly closed by the remote host" ,it may cause by your tsl setting or connection is not tsl/ssl.
        if something wrong and error message isn't useful, you can get more information in wireshark or other network tools.

        3. set Authorization in header and path for every wss requesting
        port for wss (websocket over ssl/tsl) is 443.
        beware that alway provide the port to 443 because paho default port is 1883 and 8883 witch is port for mqtt over ssl not for wss.

        4. initialize callback pools and essential callback function

        5. connect to broker and then client is ready for subscribing to any topics. well done!

        Parameters
        ----------
        host : str
            mqtt server address ,for example "mqtt.host.com"  or "192.168.1.1"
            provide only host name NOT INCLUDE port number, url path or even prefix like "mqtt://" or "tcp://" or "wss://"
            because they will be added in another place
        path : str
            mqtt path
        token_type : Literal["Bearer", "Basic"]
            token_type was return along together with access_token once you logged in
        token : str
            access token,for example "eyJhbGc..."
        client_id : str, optional
            client key for wss it can be a unique string or some unique meaning string as you want if blank it will be auto generated by library, by default ""

        Note
        ----
        host and token can get from settrade api ex (GET)https://open-api.settrade.com/api/dispatcher/v3/<broker_id>/token

        """
        self.callback_pool: Dict[str, List[Callback]] = dict()
        self.subscribed_topics: Set[str] = set()

        client = mqtt.Client(
            clean_session=True,
            client_id=client_id,
            transport="websockets",
        )

        client.tls_set()

        client.ws_set_options(
            headers={
                "Authorization": f"{token_type} {token}",
            },
            path=path,
        )

        client.on_message = self._create_callback_pool("on_message")
        client.on_connect = self._create_callback_pool("on_connect")
        client.on_disconnect = self._create_callback_pool("on_disconnect")
        client.on_connect_fail = self._create_callback_pool("on_connect_fail")
        client.on_subscribe = self._create_callback_pool("on_subscribe")
        client.on_unsubscribe = self._create_callback_pool("on_unsubscribe")
        client.on_log = self._create_callback_pool("on_log")
        client.on_publish = self._create_callback_pool("on_publish")

        self._add_base_callback()

        client.connect(host=host, port=443, keepalive=30)

        self.client = client

        self._add_base_topic(["$sys/u/_broker/_uref/error/subscribe"])

    def _add_base_callback(self):
        """Add essential callback functions to monitor and handle connection status"""

        def on_disconnect(client, userdata, rc):
            self.is_connected = False
            if rc != 0:
                self.inject_message(None, f"disconnected with rc : {rc}")

        self.add_callback("on_disconnect", "__base__", on_disconnect)

    def _add_base_topic(self, base_topics: List[str]):
        """Add essential topics to monitor and handle connection status, for example system topic

        Parameters
        ----------
        base_topics : List[str]
            Found in API documentation
        """
        for topic in base_topics:
            self.subscribe(topic)

        self.base_topic_count = len(base_topics)

    def _get_callback_pool(self, callback_type: CallbackType) -> List[Callback]:
        self.lock_callback_pool.acquire()
        return self.callback_pool[callback_type]

    def _execute_callback(self, c: Callable, *args):
        threading.Thread(target=c, args=args).start()

    def _create_callback_pool(self, callback_type: CallbackType) -> Callable:
        """Create callback function that execute all callback in callback list"""
        self.callback_pool[callback_type] = list()

        if callback_type == "on_message":

            def message_pool(client, user_data, msg):
                """in message_pool, there is system topic name $sys/u/_broker/_uref/error/subscribe
                this topic return error message of other subscribe topic, so we need to inject it to main topic
                """
                callback_pool = self._get_callback_pool(callback_type)
                for c in callback_pool:
                    if msg.topic == c.topic:
                        self._execute_callback(c.callback, client, user_data, msg)
                    if msg.topic == "$sys/u/_broker/_uref/error/subscribe":
                        err = json.loads(msg.payload.decode())
                        if c.topic == err["rejectSubscriptions"][0]["topicFilter"]:
                            self._execute_callback(c.callback, client, user_data, msg)

                self.lock_callback_pool.release()

            return message_pool
        else:

            def callback_pool(*args):
                for c in self.callback_pool[callback_type]:
                    c.callback(*args)

            return callback_pool

    def subscribe(self, topic: str):
        """Subscribe only one time per topic

        Parameters
        ----------
        topic : str
            Found in API documentation
        """
        if topic not in self.subscribed_topics:
            self.client.subscribe(topic)
            self.subscribed_topics.add(topic)

    def unsubscribe(self, topic: str):
        """Unsubscribe and remove this topic from subscribed_topics

        Parameters
        ----------
        topic : str
            Found in API documentation
        """
        if topic in self.subscribed_topics:
            self.client.unsubscribe(topic)
            self.subscribed_topics.remove(topic)

    def start(self):
        """Start connection only once"""
        if not self.is_connected:
            self.is_connected = True
            self.client.loop_start()

    def stop(self):
        """Stop and disconnect session"""
        self.is_connected = False
        self.client.loop_stop()
        self.client.disconnect()

    def stop_if_no_sub(self):
        if len(self.subscribed_topics) == self.base_topic_count:
            self.stop()

    def add_callback(self, callback_type: CallbackType, topic: str, callback: Callable):
        """Add callback to callback pool"""
        assert callback_type in CALLBACK_TYPE_LIST

        self.callback_pool[callback_type].append(
            Callback(topic, callback, callback_type)
        )

    def remove_callback(
        self,
        callback_type: CallbackType,
        topic: str,
        callback: Optional[Callable] = None,
    ):
        """Remove callback from callback pool

        auto unsub, if current topic doesn't have any on_message left in pool
        """

        assert callback_type in CALLBACK_TYPE_LIST

        callback_pool = self._get_callback_pool(callback_type)

        is_message_callback_left = False

        for i in range(len(callback_pool) - 1, -1, -1):
            c = callback_pool[i]
            if c.topic == topic and callback == None:
                del callback_pool[i]
            elif c.topic == topic and id(c.callback) == id(callback):
                del callback_pool[i]
            elif callback_type == "on_message" and c.topic == topic:
                is_message_callback_left = True
            elif callback_type != "on_message":
                is_message_callback_left = True

        if is_message_callback_left == False:
            self.unsubscribe(topic)

        self.lock_callback_pool.release()

    def inject_message(self, topic: Optional[str], message: str):
        """Inject message to specific topic or all topics"""
        callback_pool = self._get_callback_pool("on_message")
        for c in callback_pool:
            if topic == None or c.topic == topic:
                msg = mqtt.MQTTMessage()
                msg.payload = message.encode()
                msg.topic = topic
                self._execute_callback(c.callback, None, None, msg)
        self.lock_callback_pool.release()


class Subscriber:
    """This is Facade class that used to manage add/remove callback and start Callbacker"""

    def __init__(self, call_backer: CallBacker, topic: str) -> None:
        self._call_backer = call_backer
        self.topic = topic
        self.callback_list: List[Callback] = list()
        self._init_session()

    def _init_session(self):
        """subscribe to topic"""
        self._call_backer.subscribe(self.topic)

    def add_callback(
        self,
        callback_type: CallbackType,
        callback: Callable,
    ):
        """Store callback in list and wait for 'start()' for adding callback to pool"""
        self.callback_list.append(Callback(self.topic, callback, callback_type))

    def stop(self):
        """Remove all callbacks of this subscriber from pool"""

        for c in self.callback_list:
            self._call_backer.remove_callback(
                callback_type=c.callback_type, topic=c.topic, callback=c.callback
            )

        self.callback_list.clear()

        self._call_backer.stop_if_no_sub()

    def start(self):
        """Add all callbacks of this subscriber to pool"""
        for c in self.callback_list:
            self._call_backer.add_callback(
                callback_type=c.callback_type, topic=c.topic, callback=c.callback
            )

        self._call_backer.start()


class RealtimeDataConnection:
    def __init__(self, context: Context):
        self._ctx = context
        self.host: Optional[str] = None
        self.token: Optional[str] = None
        self.version = "v3"
        self.dispatcher_path = f"/api/dispatcher/{self.version}/{self._ctx.broker_id}"
        self.call_backer: Optional[CallBacker] = None

    @property
    def token_uri(self):
        return f"{self._ctx.base_url}{self.dispatcher_path}/token"

    def fetch_host_token(self):
        """
        request to settrade-dispatcher to get wss host and wss token
        """
        option = Option(
            "GET",
            self.token_uri,
        )
        response = self._ctx.dispatch(option)
        data = response.json()
        self.host = data["hosts"][0]
        self.token = data["token"]

    def get_call_backer(self, client_id: str = "") -> CallBacker:
        """Create new callbacker if it is not exist

        Parameters
        ----------
        client_id : str, optional
            an any unique string value, by default ""

        Returns
        -------
        CallBacker
            CallBacker that connected to broker
        """

        if self.host == None or self.token == None:
            self.fetch_host_token()

        if self.call_backer == None:
            self.call_backer = CallBacker(
                host=self.host,
                path=f"{self.dispatcher_path}/mqtt",
                token_type=self._ctx.token_type,
                token=self.token,
                client_id=client_id,
            )

        return self.call_backer

    def create_subscriber(
        self,
        topic: str,
        client_id: str = "",
    ) -> Subscriber:
        return Subscriber(self.get_call_backer(client_id), topic)

    def create_msg_hdlr(self, schema: betterproto.Message):
        def msg_hdlr(msg: Union[bytes, str]):
            if type(msg) is str:
                return mqtt_to_message(False, msg)
            if type(msg) is bytes:
                try:
                    if not msg.startswith(b"\n"):
                        raise Exception("invalid binary")
                    return mqtt_to_message(True, schema.parse(msg).to_dict())
                except:
                    msg_str = msg.decode()
                    if "rejectSubscriptions" in msg_str:
                        return mqtt_to_message(False, json.loads(msg))
                    return mqtt_to_message(False, msg_str)
            return mqtt_to_message(False, "something went wrong")

        return msg_hdlr

    def create_on_message(
        self,
        on_message: Callable,
        schema: Optional[betterproto.Message] = None,
        args: tuple = (),
        kwargs: dict = {},
    ):
        def callback(client: mqtt.Client, user_data: Any, msg):
            if schema == None:
                on_message(msg.payload, *args, **kwargs)
            else:
                msg_hdlr = self.create_msg_hdlr(schema)
                on_message(msg_hdlr(msg.payload), *args, **kwargs)

        return callback

    def subscribe_bid_offer(
        self,
        symbol: str,
        on_message: Callable,
        args: tuple = (),
        kwargs: dict = {},
    ) -> Subscriber:
        s = self.create_subscriber(f"proto/topic/bidoffer{self.version}/{symbol}")
        s.add_callback(
            "on_message",
            self.create_on_message(on_message, BidOfferV3(), args, kwargs),
        )
        return s

    def subscribe_derivatives_order(
        self,
        account_no: str,
        on_message: Callable,
        args: tuple = (),
        kwargs: dict = {},
    ) -> Subscriber:
        s = self.create_subscriber(
            f"proto/ua/_broker/{account_no}/_front/orderdv{self.version}"
        )
        s.add_callback(
            "on_message",
            self.create_on_message(on_message, OrderDerivV3(), args, kwargs),
        )
        return s

    def subscribe_equity_order(
        self,
        account_no: str,
        on_message: Callable,
        args: tuple = (),
        kwargs: dict = {},
    ) -> Subscriber:
        s = self.create_subscriber(
            f"proto/ua/_broker/{account_no}/_front/ordereq{self.version}"
        )
        s.add_callback(
            "on_message",
            self.create_on_message(on_message, OrderEquityV3(), args, kwargs),
        )
        return s

    def subscribe_price_info(
        self,
        symbol: str,
        on_message: Callable,
        args: tuple = (),
        kwargs: dict = {},
    ) -> Subscriber:
        s = self.create_subscriber(f"proto/topic/info{self.version}/{symbol}")
        s.add_callback(
            "on_message", self.create_on_message(on_message, InfoV3(), args, kwargs)
        )
        return s

    def subscribe_derivatives_exchange_info(
        self,
        market: str,
        on_message: Callable,
        args: tuple = (),
        kwargs: dict = {},
    ):
        s = self.create_subscriber(f"proto/topic/exchinfodv{self.version}/{market}")
        s.add_callback(
            "on_message",
            self.create_on_message(on_message, DerivExchangeInfoV3(), args, kwargs),
        )
        return s

    def subscribe_equity_exchange_info(
        self,
        market: str,
        on_message: Callable,
        args: tuple = (),
        kwargs: dict = {},
    ):
        s = self.create_subscriber(f"proto/topic/exchinfoeq{self.version}/{market}")
        s.add_callback(
            "on_message",
            self.create_on_message(on_message, EquityExchangeInfoV3(), args, kwargs),
        )
        return s

    def subscribe_error(
        self,
        on_message: Optional[Callable] = None,
        args: tuple = (),
        kwargs: dict = {},
    ):
        s = self.create_subscriber(f"$sys/u/_broker/_uref/error/subscribe")
        if on_message != None:
            s.add_callback(
                "on_message", self.create_on_message(on_message, None, args, kwargs)
            )
        return s

    def wait(self):
        c = self.get_call_backer()
        if c.client._thread != None:
            c.client._thread.join()

    def stop(self):
        if self.call_backer != None:
            self.call_backer.stop()
            del self.call_backer
