import click
import socket
import time
from datetime import datetime
import threading
import ssl
import logging
from typing import Optional
import websocket
from concurrent.futures import Future, TimeoutError
from websocket._abnf import ABNF
from elephant_socks5.protocol import bytes_to_frame, OP_CONTROL, OP_DATA, JRPCResponse, SessionRequest, Hello, Frame, TerminationRequest
from elephant_socks5.utils import LengthFieldBasedFrameDecoder, chunk_list, sneaky, socket_description, has_format_placeholders, Metric, my_ip
from elephant_socks5.version import __version__
from click import secho
import json
import string
from py_netty import Bootstrap, ServerBootstrap, ChannelHandlerAdapter, EventLoopGroup
from attrs import define, field
from itertools import cycle, count

logger = logging.getLogger("elephant-socks5")

URL = "ws[s]://localhost:4443/elephant/ws"

decoder = LengthFieldBasedFrameDecoder()

_enable_reverse_proxy = False
_quiet = False
_trace_data = False
_session_request_timeout = 3
_no_color = False

_proxy_ip = None
_proxy_port = None

_tunnels = []                   # list of Tunnel
tunnels = None                  # cycle object
_counter = count(start=0, step=1)
_lock = threading.Lock()
_connected_counter = set()

ACCEPTOR = EventLoopGroup(1, 'Acceptor')
WORKER = EventLoopGroup(1, 'Worker')


def _actives():
    return len([t for t in _tunnels if t.check_connected()])


def log_print(msg, *args, fg=None, underline=None, level=logging.INFO, force_print=False, **kwargs):
    logger.log(level, msg, *args, **kwargs)
    if not _quiet or force_print:
        message = msg
        if has_format_placeholders(msg):
            try:
                message = string.Formatter().format(msg, *args, **kwargs)
            except Exception:
                pass
            if has_format_placeholders(message):
                try:
                    message = message % args
                except Exception:
                    pass
        else:
            message = ' '.join(map(str, (msg, *args)))

        if _no_color:
            fg = None
            underline = None
        secho(message, fg=fg, underline=underline)


@define(slots=False)
class Tunnel:

    url: str = field()
    hello_params: dict = field(factory=dict)

    def __attrs_post_init__(self):
        self._clients = {}      # session_id -> ctx
        self._responseFutures = {}  # id -> (method, Future(jrpc))
        self._thread = None     # WebSocket thread
        self._localport = -1
        self._count = next(_counter)  # sequence
        self.hello_params.update({
            'version': __version__,
            'seq': self._count
        })
        self._session_create_metric = None
        self._ws = None

    def create_session_create_metric(self):
        self._session_create_metric = Metric.of(f"session-create-rtt @{self._count}")

    @sneaky()
    def start(self):
        self._ws = websocket.WebSocketApp(
            self.url,
            on_open=self._on_open,
            on_message=self._on_message,
            on_error=self._on_error,
            on_close=self._on_close
        )

        # ws.run_forever(dispatcher=rel, reconnect=5)
        log_print(f"Starting Tunnel#{self._count} (URL: {self.url}) ...", underline=True, force_print=True)
        self._ws.run_forever(reconnect=10, sslopt={"cert_reqs": ssl.CERT_NONE}, ping_interval=60 * 3, ping_timeout=30)
        # rel.signal(2, rel.abort)
        # rel.dispatch()

    def check_connected(self, throw=False):
        if not self._ws:
            if throw:
                raise Exception("Tunnel not ready (Init Connecting))")
            return False
        try:
            _ = self._ws.sock.sock.getpeername()
            connected = self._ws.sock.connected
        except Exception:
            if throw:
                raise Exception(">>Tunnel Disconnected (Reconnecting)")
            return False
        if not connected:
            if throw:
                raise Exception(">>>Tunnel Disconnected (Reconnecting)")
            return False
        return True

    def _send_frame(self, frame: Frame, method: str = None) -> None:
        self.check_connected(throw=True)
        self._trace(frame, True, method)
        self._ws.send(frame.to_bytes(), ABNF.OPCODE_BINARY)

    def _handle_frame(self, ws, frame_bytes):
        frame = bytes_to_frame(frame_bytes)
        self._trace(frame, send=False)
        if frame.op_type == OP_CONTROL:
            payload = frame.payload
            jrpc = json.loads(payload.decode('utf-8'))
            jrpc_id = jrpc['id']

            if jrpc_id in self._responseFutures:
                method, future = self._responseFutures.get(jrpc_id)
                del self._responseFutures[jrpc_id]
            else:                   # request
                method = jrpc.get('method')

            if method == 'echo-request':
                self._send_frame(JRPCResponse.of(jrpc_id).to_frame(), 'echo-response')
            elif method == 'termination-response':
                pass
            elif method == 'termination-request':
                try:
                    session_id = jrpc['params']['session-id']
                    if session_id in self._clients:
                        ctx = self._clients[session_id]
                        ctx.close()
                        self.remove_session(session_id)
                finally:
                    self._send_frame(JRPCResponse.of(jrpc_id).to_frame(), 'termination-response')
            elif method == 'session-request-response':
                future.set_result(jrpc)
            elif method == 'agent-hello-ack':
                pass
            elif method == 'session-request':
                if not _enable_reverse_proxy:
                    self._send_frame(JRPCResponse.of(jrpc_id, error={
                        'reason': "reverse proxy is not enabled"
                    }).to_frame(), 'session-request-response')
                ip = jrpc['params']['ip']
                port = jrpc['params']['port']
                session_id = jrpc['params']['session-id']
                ok = self._create_reverse_proxy(session_id, ip, port)
                if ok:
                    self._send_frame(JRPCResponse.of(jrpc_id).to_frame(), 'session-request-response')
                else:
                    self._send_frame(JRPCResponse.of(jrpc_id, error={
                        'reason': f"fail to connect {ip}:{port}/session:{session_id}"
                    }).to_frame(), 'session-request-response')
            else:
                log_print(f"Unknown method: {method}", fg='red')
        elif frame.op_type == OP_DATA:
            session_id = frame.session_id
            payload = frame.payload
            if session_id in self._clients:
                ctx = self._clients[session_id]
                ctx.write(payload)

    def _create_reverse_proxy(self, session_id: int, ip: str, port: int) -> bool:
        try:
            b = Bootstrap(
                eventloop_group=WORKER,
                handler_initializer=lambda: ReverseProxyChannelHandler(self, session_id)
            )
            channel = b.connect(ip, port).sync().channel()
            if not channel.is_active():
                log_print(f"[Reverse]Failed to connect to {ip}:{port}", fg='red', level=logging.ERROR)
                return False
        except Exception as e:
            log_print(f"[Reverse]Failed to connect to {ip}:{port} ({e})", fg='red', level=logging.ERROR)
            return False
        self._clients[session_id] = channel.context()
        return True

    def remove_session(self, session_id: int) -> None:
        self._clients.pop(session_id, None)

    def send_termination_request(self, session_id: int) -> None:
        assert self._thread is not None and threading.current_thread() != self._thread
        if session_id not in self._clients:
            return
        tr = TerminationRequest.of(session_id)
        self._responseFutures[tr.id] = ('termination-response', Future())
        self._send_frame(tr.to_frame())

    def send_session_request(self, ctx) -> int:  # return session_id
        """This method is called by the client (Netty thread) to request a new session."""
        assert self._thread is not None and threading.current_thread() != self._thread
        if _proxy_ip and _proxy_port:
            sr = SessionRequest.of(_proxy_ip, _proxy_port)
        else:
            sr = SessionRequest.socks5()
        future = Future()
        self._responseFutures[sr.id] = ('session-request-response', future)
        s = time.perf_counter()
        self._send_frame(sr.to_frame())
        try:
            jrpc = future.result(_session_request_timeout)  # wait for response
            e = time.perf_counter()
            self._session_create_metric.record(e - s)
        except TimeoutError:
            raise Exception(f"Session request timeout [{self._session_create_metric}]")
        session_id = jrpc['result']['session-id']
        self._clients[session_id] = ctx
        return session_id

    @sneaky()
    def _on_open(self, ws):
        with _lock:
            log_print(f"[on_open {_actives()}/{len(_tunnels)}] Opened connection @{self._count} {socket_description(ws.sock.sock)}", fg='bright_blue', force_print=True)

        self._localport = ws.sock.sock.getsockname()[1]
        self._thread = threading.current_thread()
        self._ws = ws
        if self._clients:
            clients_count = len(self._clients)
            log_print(f"Closing {clients_count} stale connections ...", fg='magenta')
        for ctx in self._clients.values():
            ctx.close()
        self._responseFutures = {}
        self._clients = {}
        self.create_session_create_metric()

        obj = Hello()
        obj.params.update(self.hello_params)
        self._responseFutures[obj.id] = ('agent-hello-ack', Future())
        self._send_frame(obj.to_frame())

    @sneaky()
    def _on_message(self, ws, bytes):
        frames = decoder.decode(bytes)
        for frame in frames:
            self._handle_frame(ws, frame)

    def _on_close(self, ws, close_status_code, close_msg):
        with _lock:
            log_print(f"[on_close {_actives()}/{len(_tunnels)}] Connection closed @{self._count}, status code: {close_status_code}, message: {close_msg} [{self._session_create_metric}]", fg='red', force_print=True)

    def _on_error(self, ws, error):
        with _lock:
            log_print(f"[on_error @{self._count} {_actives()}/{len(_tunnels)}] {error} [{self._session_create_metric}]", fg='red', level=logging.ERROR, force_print=True)

    def _trace(self, frame: Frame, send=True, method: str = None):
        if _quiet:
            return
        tunnel_identifier = f"#{self._count}:{self._localport}".center(10)
        direction = '>>>' if send else '<<<'
        fg = 'yellow' if not send else None
        now = f"{str(datetime.now()):<20}"
        if frame.op_type == OP_CONTROL:
            jrpc = json.loads(frame.payload.decode('utf-8'))
            jrpc_id = jrpc['id']
            method = method or jrpc.get('method')
            if not method:
                if jrpc_id in self._responseFutures:
                    method, _ = self._responseFutures.get(jrpc_id)
                else:
                    method = '?'

            if 'agent-hello' == method:
                msg0 = jrpc.get('params')
            elif 'agent-hello-ack' == method:
                msg0 = jrpc.get('result')
            elif 'response' in method:
                e = jrpc.get('error')
                msg0 = e or jrpc.get('result')
                if e:
                    fg = 'red'
            elif 'request' in method:
                msg0 = jrpc.get('params', jrpc)
            else:
                msg0 = jrpc

            if isinstance(msg0, dict):
                msg0 = json.dumps(msg0)

            session_count = len(self._clients)
            if method == 'session-request-response' and not send:
                session_count += 1

            msg = f"{now} | {tunnel_identifier} | {direction} | {method:<25} | {jrpc_id} | sessions: {session_count}"

            if logger.isEnabledFor(logging.DEBUG):
                msg += f" | futures: {len(self._responseFutures)}"
            msg += f" | {msg0}"
        elif frame.op_type == OP_DATA:
            if not _trace_data:
                return
            session_id = frame.session_id
            msg = f"{now} | {tunnel_identifier} | {direction} | [{'OP_DATA: ' + str(session_id):<25}] <<{len(frame.payload)} bytes>>"
            ctx = self._clients.get(session_id)
            if ctx:
                msg += f" {ctx.channel()}"
            if not send and not ctx:
                fg = 'red'
        else:
            msg = f"{now} | {tunnel_identifier} | {direction} | Are you kidding me?"

        log_print(msg, fg=fg, level=logging.DEBUG)


class ReverseProxyChannelHandler(ChannelHandlerAdapter):

    def __init__(self, tunnel: Tunnel, session_id: int):
        self._session_id = session_id
        self._tunnel = tunnel
        self._tunnel_seq = tunnel._count

    def exception_caught(self, ctx, exception):
        logger.error("[Reverse][Exception Caught #%d @%d] %s : %s",
                     self._session_id, self._tunnel_seq, ctx.channel(), str(exception), exc_info=exception)
        click.secho(f"[Reverse][Exception Caught #{self._session_id} @{self._tunnel_seq}] {ctx.channel()} : {str(exception)}", fg='red')
        ctx.close()

    def channel_active(self, ctx):
        log_print(f"[Reverse][channel_active #{self._session_id} @{self._tunnel_seq}] {ctx.channel()}", fg='green', level=logging.DEBUG)

    def channel_read(self, ctx, bytebuf):
        for sub_bytebuf in chunk_list(bytebuf, 1024):
            data_frame = Frame(
                op_type=OP_DATA,
                session_id=self._session_id,
                payload=sub_bytebuf,
            )
            self._tunnel._send_frame(data_frame)

    def channel_inactive(self, ctx):
        log_print(f"[Reverse][channel_inactive #{self._session_id} @{self._tunnel_seq}] {ctx.channel()}", fg='bright_black', level=logging.DEBUG)
        self._tunnel.remove_session(self._session_id)
        self._tunnel.send_termination_request(self._session_id)


class ProxyChannelHandler(ChannelHandlerAdapter):

    def __init__(self):
        self._session_id = None
        self._tunnel = self._next_valid_tunnel()
        self._tunnel_seq = self._tunnel._count if self._tunnel else '?'

    @staticmethod
    def _next_valid_tunnel() -> Optional[Tunnel]:
        for _ in range(len(_tunnels)):
            t = next(tunnels)
            if t.check_connected():
                return t

    def exception_caught(self, ctx, exception):
        logger.error("[Exception Caught #%s @%s] %s : %s", self._session_id or '?', self._tunnel_seq, ctx.channel(), str(exception), exc_info=exception)
        click.secho(f"[Exception Caught #{self._session_id or '?'} @{self._tunnel_seq}] {ctx.channel()} : {str(exception)}", fg='red')
        ctx.close()

    def channel_read(self, ctx, bytebuf):
        if not self._session_id:
            ctx.close()         # failed to get seesion_id and there is incoming data here
            return
        for sub_bytebuf in chunk_list(bytebuf, 1024):
            data_frame = Frame(
                op_type=OP_DATA,
                session_id=self._session_id,
                payload=sub_bytebuf,
            )
            self._tunnel._send_frame(data_frame)

    def channel_active(self, ctx):
        log_print(f"[channel_active @{self._tunnel_seq}] {ctx.channel()}", fg='green', level=logging.DEBUG)
        if self._tunnel is None:
            log_print(f"{ctx.channel()} All tunnels disconnected!!", fg='red', level=logging.ERROR)
            ctx.close()
            return
        self._session_id = self._tunnel.send_session_request(ctx)
        if self._session_id < 0:
            log_print(f"Failed to establish session with {ctx.channel()}", fg='red', level=logging.ERROR)
            ctx.close()
            return

    def channel_inactive(self, ctx):
        log_print(f"[channel_inactive #{self._session_id or '?'} @{self._tunnel_seq}] {ctx.channel()}", fg='bright_black', level=logging.DEBUG)
        if self._session_id:
            try:
                self._tunnel.send_termination_request(self._session_id)
            finally:
                self._tunnel.remove_session(self._session_id)


def _config_logging():
    from logging.handlers import RotatingFileHandler

    logging.basicConfig(
        handlers=[
            RotatingFileHandler(
                filename="elephant-client.log",
                maxBytes=10 * 1024 * 1024,  # 10M
                backupCount=10,
            ),
            # logging.StreamHandler(),  # default to stderr
        ],
        level=logging.INFO,
        format='%(asctime)s.%(msecs)03d - %(name)s - %(levelname)s - %(threadName)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    logger.setLevel(logging.INFO)


@click.command(short_help="Elephant SOCKS5 client", context_settings=dict(help_option_names=['-h', '--help'], max_content_width=120))
@click.option('--port', '-p', default=1080, help="Local port to bind", show_default=True, type=int)
@click.option('--global', '-g', 'global_', is_flag=True, help="Listen on all interfaces")
@click.option('--server', '-s', 'urls', help=f"Elephant tunnel server URLs (like: {URL})", type=str, required=True)
@click.option('--alias', '-a', help="Alias for the client", type=str)
@click.option('--quiet', '-q', is_flag=True, help="Quiet mode")
@click.option('--enable-reverse-proxy', '-r', 'enable_reverse_proxy', is_flag=True, help="Enable reverse proxy")
@click.option('--reverse-proxy-only', '-rpo', 'reverse_proxy_only', is_flag=True, help="No SOCKS5 server, only for reverse proxy")
@click.option('--reverse-ip', help="Reverse proxy IP", type=str)
@click.option('--reverse-port', help="Reverse proxy port", type=int, default=-1, show_default=True)
@click.option('--no-reverse-global', 'no_reverse_global', is_flag=True, help="Reverse proxy listen on localhost")
@click.option('--log-record', '-l', 'save_log', is_flag=True, help="Save log to file (elephant-client.log)")
@click.option('--request-timeout', '-t', 'session_request_timeout', default=3, help="Session request timeout (seconds)", show_default=True, type=int)
@click.option('--no-color', is_flag=True, help="Disable color output")
@click.option('--verbose', '-v', 'verbose', is_flag=True, help="Verbose mode")
@click.option('--tunnels', '-n', 'tunnel_count', default=1, help="Number of tunnels to achieve load balance", show_default=True, type=click.IntRange(1, clamp=True))
@click.option('--proxy-ip', help="Proxy IP", type=str)
@click.option('--proxy-port', help="Proxy port", type=int, default=-1, show_default=True)
@click.version_option(version=__version__, prog_name="Elephant SOCKS5 Client")
def _cli(
        port, urls, alias, tunnel_count,
        quiet, save_log, session_request_timeout, no_color,
        proxy_ip, proxy_port, global_,
        enable_reverse_proxy, reverse_ip, reverse_port, no_reverse_global,
        verbose, reverse_proxy_only,
):
    global _quiet
    global _session_request_timeout
    global _no_color
    global _tunnels
    global tunnels              # cycle list of tunnels
    global _proxy_ip
    global _proxy_port
    global _enable_reverse_proxy

    if proxy_ip:
        _proxy_ip = proxy_ip
        assert 0 < proxy_port < 65536, "Invalid proxy port"
        _proxy_port = proxy_port

    _no_color = no_color
    _quiet = quiet
    _session_request_timeout = session_request_timeout
    _enable_reverse_proxy = enable_reverse_proxy or reverse_proxy_only or (reverse_ip and reverse_port > 0)

    hello_params = {}
    hello_params['alias'] = alias or socket.gethostname()
    hello_params['myip'] = my_ip()
    if reverse_ip and reverse_port > 0:
        hello_params['reverseHost'] = reverse_ip
        hello_params['reversePort'] = reverse_port
        hello_params['reverseGlobal'] = not no_reverse_global

    if save_log or quiet:
        _config_logging()

    if verbose:
        logger.setLevel(logging.DEBUG)

    urls = set(urls.split(','))
    tunnel_count = max(tunnel_count, len(urls))
    urls = cycle(urls)
    _tunnels = [Tunnel(next(urls), hello_params.copy()) for _ in range(tunnel_count)]
    tunnels = cycle(_tunnels)
    for t in _tunnels:
        wst = threading.Thread(target=t.start)
        wst.daemon = True
        wst.start()

    log_print(f"Waiting for all {len(_tunnels)} tunnels to be ready ...", underline=True, force_print=True)
    count = 10
    while not all(t.check_connected() for t in _tunnels):
        count -= 1
        time.sleep(0.5)
        if count <= 0:
            log_print("Timeout waiting for all tunnels to be ready, abort", fg='red', level=logging.ERROR, force_print=True)
            return

    if reverse_proxy_only:
        log_print("Reverse proxy only mode!!", underline=True, force_print=True)
        while True:
            time.sleep(86400)
        return

    log_print(f"Proxy server started and listening on port {port} ...", fg='green', underline=True, force_print=True)
    sb = ServerBootstrap(
        parant_group=ACCEPTOR,
        child_group=WORKER,
        child_handler_initializer=ProxyChannelHandler
    )

    sb.bind(address='0.0.0.0' if global_ else '127.0.0.1', port=port).close_future().sync()


def _run():
    _cli()
