import logging
import os
import socket
import traceback

import ray

import pyjava.utils as utils
from pyjava.api.mlsql import RayContext
from pyjava.rayfix import RayWrapper
from pyjava.serializers import \
    write_with_length, \
    write_int, read_int, \
    SpecialLengths, ArrowStreamPandasSerializer


class SocketNotBindException(Exception):
    def __init__(self, message):
        Exception.__init__(self)
        self.message = message


class DataServerWithId(object):
    def __init__(self, host, port, server_id):
        self.host = host
        self.port = port
        self.server_id = server_id


class OnceServer(object):
    def __init__(self, host, port, timezone):
        self.host = host
        self.port = port
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.settimeout(5 * 60)
        self.out_ser = ArrowStreamPandasSerializer(timezone, False, None)
        self.is_bind = False
        self.is_dev = utils.is_dev()

    def bind(self):
        try:
            self.socket.bind((self.host, self.port))
            self.is_bind = True
            self.socket.listen(1)
        except Exception:
            print(traceback.format_exc())

        return self.socket.getsockname()

    def close(self):
        self.socket.close()

    def serve(self, data):
        from pyjava.api.mlsql import PythonContext
        if not self.is_bind:
            raise SocketNotBindException(
                "Please invoke server.bind() before invoke server.serve")
        conn, addr = self.socket.accept()
        sockfile = conn.makefile("rwb", int(
            os.environ.get("BUFFER_SIZE", 65536)))
        infile = sockfile  # os.fdopen(os.dup(conn.fileno()), "rb", 65536)
        out = sockfile  # os.fdopen(os.dup(conn.fileno()), "wb", 65536)
        try:
            write_int(SpecialLengths.START_ARROW_STREAM, out)
            out_data = ([df[name] for name in df] for df in
                        PythonContext.build_chunk_result(data, 1024))
            self.out_ser.dump_stream(out_data, out)

            write_int(SpecialLengths.END_OF_DATA_SECTION, out)
            write_int(SpecialLengths.END_OF_STREAM, out)
            out.flush()
            if self.is_dev:
                print("all data  in ray task have been consumed.")
            read_int(infile)
        except Exception:
            try:
                write_int(SpecialLengths.ARROW_STREAM_CRASH, out)
                ex = traceback.format_exc()
                print(ex)
                write_int(SpecialLengths.PYTHON_EXCEPTION_THROWN, out)
                write_with_length(ex.encode("utf-8"), out)
                out.flush()
                read_int(infile)
            except IOError:
                # JVM close the socket
                pass
            except Exception:
                # Write the error to stderr if it happened while serializing
                print("Py worker failed with exception:")
                print(traceback.format_exc())
                pass

        conn.close()


@ray.remote
class RayDataServer(object):

    def __init__(self, server_id, java_server, port=0, timezone="Asia/Harbin"):
        self.server = OnceServer(
            RayWrapper().get_address(), port, java_server.timezone)
        try:
            (rel_host, rel_port) = self.server.bind()
        except Exception as e:
            print(traceback.format_exc())
            raise e

        self.host = rel_host
        self.port = rel_port
        self.timezone = timezone
        self.server_id = server_id
        self.java_server = java_server
        self.is_dev = utils.is_dev()

    def serve(self, func_for_row=None, func_for_rows=None):
        try:
            if func_for_row is not None:
                data = (func_for_row(item)
                        for item in RayContext.fetch_once_as_rows(self.java_server))
            elif func_for_rows is not None:
                data = func_for_rows(
                    RayContext.fetch_once_as_rows(self.java_server))
            self.server.serve(data)
        except Exception as e:
            logging.error(f"Fail to processing data in  Ray Data Server {self.host}:{self.port}")
            raise e
        finally:
            self.close()

    def close(self):
        try:
            self.server.close()
        except Exception:
            print(traceback.format_exc())
        ray.actor.exit_actor()

    def connect_info(self):
        return DataServerWithId(self.host, self.port, self.server_id)
