import logging
import traceback

from thrift.protocol import TCompactProtocol
from thrift.transport import TSocket, TTransport

from rtp_llm.aios.kmonitor.python_client.flume import ThriftSourceProtocol

logger = logging.getLogger("flume")


class _Transport(object):
    def __init__(self, thrift_host, thrift_port, timeout=None, unix_socket=None):
        self.thrift_host = thrift_host
        self.thrift_port = thrift_port
        self.timeout = timeout
        self.unix_socket = unix_socket

        self._socket = TSocket.TSocket(
            self.thrift_host, self.thrift_port, self.unix_socket
        )
        self._transport_factory = TTransport.TFramedTransportFactory()
        self._transport = self._transport_factory.getTransport(self._socket)

    def connect(self):
        try:
            if self.timeout:
                self._socket.setTimeout(self.timeout)
            if not self.is_open():
                self._transport = self._transport_factory.getTransport(self._socket)
                self._transport.open()
        except Exception as e:
            logger.warn("connect to flume exception:%s", e)
            logger.warn(traceback.format_exc())
            self.close()

    def reconnect(self):
        self.close()
        self.connect()

    def is_open(self):
        return self._transport.isOpen()

    def get_transport(self):
        return self._transport

    def close(self):
        self._transport.close()


class FlumeClient(object):
    def __init__(self, thrift_host, thrift_port, timeout=10000, unix_socket=None):
        self._transObj = _Transport(
            thrift_host, thrift_port, timeout=timeout, unix_socket=unix_socket
        )
        self._protocol = TCompactProtocol.TCompactProtocol(
            trans=self._transObj.get_transport()
        )
        self.client = ThriftSourceProtocol.Client(
            iprot=self._protocol, oprot=self._protocol
        )
        self._transObj.connect()

    def send(self, event):
        try:
            self.client.send_append(event)
        except Exception as e:
            logger.warn("send to flume exception:%s", e)
            logger.warn(traceback.format_exc())
            self._transObj.reconnect()

    def send_batch(self, events):
        try:
            self.client.send_appendBatch(events)
        except Exception as e:
            logger.warn("send batch to flume exception:%s", e)
            logger.warn(traceback.format_exc())
            self._transObj.reconnect()

    def close(self):
        self._transObj.close()
