# -*- coding: utf-8 -*-
#
# This file is part of the bliss project
#
# Copyright (c) 2015-2023 Beamline Control Unit, ESRF
# Distributed under the GNU LGPLv3. See LICENSE for more info.
import os
import sys

import errno
import gevent
from gevent import socket, select, event

from bliss.comm.tcp import Tcp
from bliss.config.conductor.client import Lock
from bliss.config.channels import Channel


def _wait_pid(read_pipe, pid):
    while True:
        r, _, _ = select.select([read_pipe], [], [])
        if r:
            out = os.read(read_pipe, 8192)
            if not out:
                os.waitpid(pid, 0)
                break


def get_proxy_channel_name(host, port):
    return "proxy:%s:%d" % (host, port)


class Proxy(object):
    TCP = 0

    def __init__(self, config):
        if "tcp" in config:
            tcp_config = config.get("tcp")
            self.external = config.get("external", False)
            if hasattr(config, "clone"):
                self._config = tcp_config.clone()
            else:
                self._config = tcp_config.copy()
            self.__eol = tcp_config.get("eol", "\n")
            self._mode = self.TCP
            cnx = Tcp(**tcp_config)
            self.name = "%s:%d" % (cnx._host, cnx._port)
        else:
            raise NotImplementedError("Proxy: Not managed yet")

        self._cnx = None
        self._join_task = None
        self._url_channel = Channel(get_proxy_channel_name(cnx._host, cnx._port))

    def close(self):
        if not self.external:
            self._url_channel.value = None
            if self._join_task is not None:
                self._join_task.join()
            self._join_task = None

    def __info__(self):
        url = self._url_channel.value
        info = f"TCP PROXY for {self.name}\n"
        if self.external:
            info += f"External process on port {url}"
        else:
            info = f"Forked process on port {url}"
        return info

    def __getattr__(self, name):
        if name.startswith("__"):
            raise AttributeError(name)

        was_connected = None
        if not name.startswith("_"):
            was_connected = self._check_connection()

        attr = getattr(self._cnx, name)
        if callable(attr) and was_connected is False:

            def wrapper_func(*args, **kwargs):
                try:
                    return attr(*args, **kwargs)
                except socket.error as e:
                    if e.errno == errno.EPIPE:
                        raise socket.error(errno.ECONNREFUSED, "Connection refused")
                else:
                    raise

            return wrapper_func
        return attr

    @property
    def _eol(self):
        return self.__eol

    def _check_connection(self):
        if self._mode == self.TCP:
            if self._cnx is None or not self._cnx._connected:
                if hasattr(self._config, "clone"):
                    local_cfg = self._config.clone()
                else:
                    local_cfg = self._config.copy()
                url = local_cfg.pop("url")
                cnx = Tcp(url, **local_cfg)
                host, port = cnx._host, cnx._port
                if not self.external:
                    proxy_url = self._fork_server(host, port)
                else:
                    proxy_url = self._connect_to_server(host, port)
                self._cnx = Tcp(proxy_url, **local_cfg)
                try:
                    self._cnx.connect()
                except ConnectionRefusedError:
                    raise ConnectionError(
                        f"Failed to connect to tcp_proxy server for {host}:{port}"
                    )
                return False

    def _connect_to_server(self, host, port):
        local_url = self._url_channel.value
        if local_url is None:
            self._cnx = None
            raise ConnectionError(
                f"No tcp_proxy server running for {host}:{port}. Restart it !!"
            )
        return local_url

    def _fork_server(self, host, port):
        with Lock(self):
            sync = event.Event()

            def port_cbk(proxy_url):
                if not proxy_url:
                    # filter default value
                    return
                sync.set()

            try:
                self._url_channel.register_callback(port_cbk)
                local_url = self._url_channel.value
                if local_url is None:
                    self._join_task = self._real_server_fork(host, port)
                    gevent.sleep(0)
                    sync.wait()
                    local_url = self._url_channel.value
                return local_url
            finally:
                self._url_channel.unregister_callback(port_cbk)

    def _real_server_fork(self, host, port):
        script_name = __file__
        read, write = os.pipe()
        pid = os.fork()
        if pid == 0:  # child
            os.dup2(write, sys.stdout.fileno())
            os.dup2(write, sys.stderr.fileno())
            os.closerange(3, write + 1)
            os.execl(
                sys.executable,
                sys.executable,
                script_name,
                "--port",
                str(port),
                "--host",
                host,
            )
            sys.exit(0)
        else:
            os.close(write)
            wait_greenlet = gevent.spawn(_wait_pid, read, pid)
            wait_greenlet.start()
            return wait_greenlet


def get_host_port_from_beacon_name(name):
    from bliss.config.static import get_config

    beacon_config = get_config()
    config = beacon_config.get_config(name)
    if config is None:
        raise ValueError(f"Cannot find beacon object {name}")
    proxy_config = config.get("tcp-proxy")
    if proxy_config is None:
        raise ValueError(f"No tcp-proxy config for object {name}")
    tcp_config = proxy_config.get("tcp")
    if tcp_config is None:
        raise ValueError(f"Missing tcp section in tcp-proxy for object {name}")
    url = tcp_config.get("url")
    if url is None:
        raise ValueError(f"Missing url in config for object {name}")
    host, port = url.split(":")
    return host, int(port)


def main():  # proxy server part
    import signal

    global pipe_write
    pipe_read, pipe_write = os.pipe()

    def _stop_server(*args):
        # import os
        os.close(pipe_write)

    signal.signal(signal.SIGINT, _stop_server)
    signal.signal(signal.SIGTERM, _stop_server)

    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--host", dest="host", type=str, help="destination host")
    parser.add_argument("--port", dest="port", type=int, help="destination port")
    parser.add_argument(
        "--beacon-name", dest="beacon_name", type=str, help="beacon object name"
    )
    _options = parser.parse_args()

    if _options.beacon_name is not None:
        (host, port) = get_host_port_from_beacon_name(_options.beacon_name)
    else:
        if _options.host is not None and _options.port is not None:
            host = _options.host
            port = _options.port
        else:
            raise ValueError("Specficy options (--host AND --port) OR (--beacon-name)")

    tcp = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    tcp.bind(("", 0))
    tcp.listen(16)
    proxy_port = tcp.getsockname()[1]
    proxy_host = socket.gethostname()

    server_url = "%s:%d" % (proxy_host, proxy_port)

    global dont_reset_channel
    dont_reset_channel = False

    def channel_cbk(value):
        global dont_reset_channel
        if value != server_url:
            dont_reset_channel = True
            try:
                os.close(pipe_write)
            except OSError:
                pass

    channel_name = get_proxy_channel_name(host, port)
    channel = Channel(channel_name, value=server_url, callback=channel_cbk)

    runFlag = True
    fd_list = [tcp, pipe_read]
    global client
    global dest
    client = None
    dest = None
    try:
        while runFlag:
            rlist, _, _ = select.select(fd_list, [], [])
            for s in rlist:
                if s == tcp:
                    accept_flag = True
                    try:
                        if dest is None:
                            dest = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                            dest.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                            dest.connect((host, port))
                            fd_list.append(dest)
                    except:
                        dest = None
                        accept_flag = False

                    if client is not None:
                        fd_list.remove(client)
                        client.close()
                        client = None

                    client, addr = tcp.accept()
                    if accept_flag:
                        client.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                        fd_list.append(client)
                    else:
                        client.close()
                        client = None

                elif s == client:
                    try:
                        raw_data = client.recv(16 * 1024)
                    except:
                        raw_data = None

                    if raw_data:
                        dest.sendall(raw_data)
                    else:
                        fd_list.remove(client)
                        client.close()
                        client = None
                elif s == dest:
                    try:
                        raw_data = dest.recv(16 * 1024)
                    except:
                        runFlag = False
                        raw_data = None

                    if raw_data:
                        client.sendall(raw_data)
                    else:
                        dest.close()
                        fd_list.remove(dest)
                        dest = None
                        fd_list.remove(client)
                        client.close()
                        client = None
                elif s == pipe_read:
                    runFlag = False
                    break
    finally:
        if dont_reset_channel is False:
            channel.value = None
            gevent.sleep(0)


if __name__ == "__main__":
    main()
