# @Coding: UTF-8
# @Time: 2024/9/20 12:29
# @Author: xieyang_ls
# @Filename: websocket.py

import time

import socket

import base64

import hashlib

from abc import ABC

from pyutils_spirit.style import draw_websocket_banner

from logging import info, basicConfig, INFO, exception

from pyutils_spirit.concurrent_thread.thread_holder import ThreadHolder

from pyutils_spirit.concurrent_thread.thread_executor import ThreadExecutor

basicConfig(level=INFO)


class WebSocket(ABC):
    __message_buffer = None

    __partial_data = None

    __socket_server: socket.socket = None

    __listener_count = None

    __buffer_size: int = None

    __websocket_server_class: type = None

    __executor: ThreadExecutor = None

    __container: ThreadHolder[str, object] = None

    __encoding: str = None

    __onopen_func: callable = None

    __onmessage_func: callable = None

    __close_signature = None

    __onclose_func: callable = None

    __onerror_func: callable = None

    @classmethod
    def __start_server(cls, host, port, other_cls):
        cls.__partial_data = {}
        cls.__message_buffer = {}
        cls.__socket_server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        cls.__listener_count = 100000
        cls.__buffer_size = 4096
        cls.__websocket_server_class = other_cls
        cls.__executor = ThreadExecutor(executor_count=cls.__listener_count + 1)
        cls.__container = ThreadHolder.constructor()
        cls.__encoding = "utf-8"
        cls.__socket_server.bind((host, port))
        cls.__socket_server.listen(cls.__listener_count)
        cls.__executor.execute(cls.__listener_connection)
        cls.__close_signature = ""
        info(f"WebSocket Server listening on ws://{host}:{port}")

    class Session:
        __socket: socket.socket = None

        __address: tuple[str, int] = None

        __buffer_size: int = None

        __is_connected: bool = None

        __headers: list[str] = None

        __start_connected_timer: float = None

        __close_connected_timer: float = None

        __ON_CLOSE_FUNC: callable = None

        __MESSAGE_BUFFER: dict = None

        __PARTIAL_DATA: dict = None

        __WEBSOCKET_CONTAINER: ThreadHolder[str, object] = None

        def __init__(self,
                     socket_client: socket.socket,
                     address: tuple[str, int],
                     onclose_func: callable,
                     message_buffer: dict,
                     partial_data: dict,
                     is_connected: bool = True,
                     buffer_size: int = 4096):
            self.__socket = socket_client
            self.__address = address
            Session.__ON_CLOSE_FUNC = onclose_func
            Session.__MESSAGE_BUFFER = message_buffer
            Session.__PARTIAL_DATA = partial_data
            self.__is_connected = is_connected
            self.__buffer_count = buffer_size
            self.__start_connected_timer = time.time()
            Session.__WEBSOCKET_CONTAINER = ThreadHolder.constructor()

        def get_socket_client(self) -> socket.socket:
            return self.__socket

        def get_is_connected(self) -> bool:
            return self.__is_connected

        def get_address(self):
            if self.__address is None:
                self.__address = self.__socket.getpeername()
            return self.__address

        def get_headers(self) -> list[str]:
            if self.__headers is None:
                socket_headers = self.__socket.recv(self.__buffer_count).decode()
                self.__headers = socket_headers.splitlines()
            return self.__headers

        def get_params(self) -> dict | None:
            if self.__headers is None:
                return None
            header = self.__headers[0]
            arguments = header.split()[1].split("?")[1]
            params = dict(param.split('=') for param in arguments.split('&'))
            return params

        def send_message(self, message: str, encoding: str = "utf-8") -> None:
            message_bytes = message.encode(encoding)
            if len(message_bytes) <= 125:
                frame = bytearray([0x81, len(message_bytes)]) + message_bytes
            elif 126 <= len(message_bytes) <= 65535:
                frame = bytearray([0x81, 126]) + len(message_bytes).to_bytes(2, 'big') + message_bytes
            else:
                frame = bytearray([0x81, 127]) + len(message_bytes).to_bytes(8, 'big') + message_bytes
            self.__socket.send(frame)

        def close(self) -> None:
            if self.__is_connected is True:
                self.__is_connected = False
                self.__close_connected_timer = time.time()
                info(f"Disconnecting Session: {self}")
                self.__socket.close()
                ws_conn_point = Session.__WEBSOCKET_CONTAINER.remove_holder(key=str(self.__address))
                if Session.__ON_CLOSE_FUNC is not None:
                    Session.__ON_CLOSE_FUNC(self=ws_conn_point, session=self)
                session_id = id(self.__socket)
                if session_id in Session.__MESSAGE_BUFFER:
                    del Session.__MESSAGE_BUFFER[session_id]
                if self.__socket in Session.__PARTIAL_DATA:
                    del Session.__PARTIAL_DATA[self.__socket]

        def __str__(self) -> str:
            if self.__is_connected:
                connected_timer: str = f"current_connected_timer={time.time() - self.__start_connected_timer}"
            else:
                connected_timer = f"total_connected_timer={self.__close_connected_timer - self.__start_connected_timer}"
            return ("{" + f"socket={self.__socket}, " +
                    f"address={self.__address}, " +
                    f"{connected_timer}" + "}")

    @classmethod
    def __handshaking(cls, session: Session, header_item):
        # WebSocket 握手
        accept_key = base64.b64encode(
            hashlib.sha1(
                (header_item.split(": ")[1] + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode()).digest()).decode()
        response = f"HTTP/1.1 101 Switching Protocols\r\n" \
                   f"Upgrade: websocket\r\n" \
                   f"Connection: Upgrade\r\n" \
                   f"Sec-WebSocket-Accept: {accept_key}\r\n\r\n"
        session.get_socket_client().send(response.encode())

    @classmethod
    def __listener_connection(cls) -> None:
        while True:
            try:
                socket_client, addr = cls.__socket_server.accept()
                session = cls.Session(socket_client=socket_client,
                                      address=addr,
                                      onclose_func=cls.__onclose_func,
                                      message_buffer=cls.__message_buffer,
                                      partial_data=cls.__partial_data)
                for header_item in session.get_headers():
                    if "Sec-WebSocket-Key" in header_item:
                        cls.__handshaking(session, header_item)
                instance = cls.__websocket_server_class()
                cls.__container.set_holder(key=str(addr), value=instance)
                if cls.__onclose_func is not None:
                    cls.__onopen_func(instance, session=session)
                cls.__executor.execute(cls.__listener_message, session=session)
                info(f"Session: {session} connection successful")
            except Exception as e:
                exception(e)

    @classmethod
    def __receive_message(cls, session: Session):
        # 接收消息
        socket_client: socket.socket = session.get_socket_client()
        session_id = id(socket_client)
        data = cls.__partial_data.get(socket_client, b'') + socket_client.recv(cls.__buffer_size)

        try:
            while len(data) >= 2:
                # 1. 解析帧头
                fin = (data[0] & 0x80) >> 7
                rsv = (data[0] & 0x70) >> 4
                opcode = data[0] & 0x0F
                masked = (data[1] & 0x80) >> 7
                payload_len = data[1] & 0x7F

                # 2. 协议合规性校验
                if rsv != 0:
                    raise ValueError("RSV位违规")
                if not masked:  # 所有客户端帧必须掩码
                    raise ValueError("未掩码数据帧")

                # 3. 解析payload长度和掩码
                mask_start = 2
                if payload_len == 126:
                    if len(data) < 4:
                        break
                    payload_len = int.from_bytes(data[2:4], 'big')
                    mask_start = 4
                elif payload_len == 127:
                    if len(data) < 10:
                        break
                    if (data[2] & 0x80) != 0:
                        raise ValueError("无效长度")
                    payload_len = int.from_bytes(data[2:10], 'big')
                    mask_start = 10

                # 4. 检查数据完整性
                if len(data) < mask_start + 4 + payload_len:
                    break

                # 5. 解码payload
                masks = data[mask_start:mask_start + 4]
                payload = data[mask_start + 4:mask_start + 4 + payload_len]
                decoded = bytes([payload[i] ^ masks[i % 4] for i in range(payload_len)])

                # 6. 处理控制帧
                if opcode in (0x8, 0x9, 0xA):
                    if not fin:
                        raise ValueError("控制帧不能分片")
                    if payload_len > 125:
                        raise ValueError("控制帧负载过长")

                    if opcode == 0x8:  # Close帧
                        close_code = 1000
                        if len(decoded) >= 2:
                            close_code = int.from_bytes(decoded[:2], 'big')
                        # 发送Close响应
                        response = bytes([0x88, 0x02]) + close_code.to_bytes(2, 'big')
                        socket_client.send(response)
                        return session.close()
                        # 清理数据

                    elif opcode == 0x9:  # Ping帧
                        # 发送Pong帧，携带解码后的数据
                        header = bytearray([0x8A])  # FIN=1, opcode=0xA
                        if len(decoded) <= 125:
                            header.append(len(decoded))
                        elif len(decoded) <= 0xFFFF:
                            header.extend([126, *len(decoded).to_bytes(2, 'big')])
                        else:
                            header.extend([127, *len(decoded).to_bytes(8, 'big')])
                        socket_client.send(bytes(header) + decoded)
                        # 继续处理剩余数据
                        data = data[mask_start + 4 + payload_len:]
                        continue
                    elif opcode == 0xA:  # Pong帧，忽略
                        data = data[mask_start + 4 + payload_len:]
                        continue

                # 7. 处理数据帧
                if session_id not in cls.__message_buffer:
                    cls.__message_buffer[session_id] = {'fragments': [], 'opcode': opcode}
                else:
                    if opcode != 0x0:
                        raise ValueError("分片连续性错误")

                cls.__message_buffer[session_id]['fragments'].append(decoded)

                if fin:
                    full_data = b''.join(cls.__message_buffer[session_id]['fragments'])
                    original_opcode = cls.__message_buffer[session_id]['opcode']
                    del cls.__message_buffer[session_id]
                    # 处理数据
                    if original_opcode == 0x1:
                        try:
                            result = full_data.decode('utf-8')
                        except UnicodeDecodeError:
                            socket_client.send(b'\x88\x02\x03\xF7')  # 1007错误
                            return session.close()
                    else:
                        result = full_data
                    data = data[mask_start + 4 + payload_len:]
                    cls.__partial_data[socket_client] = data
                    return result
                else:
                    data = data[mask_start + 4 + payload_len:]
        except (ConnectionResetError, ValueError) as e:
            exception(f"协议错误: {e}")
            session.close()
        finally:
            cls.__partial_data[socket_client] = data
        return None

    @classmethod
    def __listener_message(cls, session: Session) -> None:
        address: tuple[str, int] = session.get_address()
        while session.get_is_connected():
            try:
                message: str = cls.__receive_message(session=session)
                if message is None:
                    continue
                if message == "ping":
                    session.send_message("pong")
                if message is cls.__close_signature:
                    session.close()
                if cls.__onmessage_func is not None:
                    instance = cls.__container.get_holder(key=str(address))
                    cls.__onmessage_func(instance, message=message)
            except Exception as e:
                exception(e)
                if cls.__onerror_func is not None:
                    instance = cls.__container.get_holder(key=str(address))
                    try:
                        cls.__onerror_func(instance, session=session, ex=e)
                    except Exception as e:
                        exception(e)

    @classmethod
    def websocket_server(cls, host: str, port: int) -> callable:

        def decorator_handler(other_cls) -> type:
            draw_websocket_banner()
            cls.__start_server(host, port, other_cls)
            other_cls.__decorator__ = "WebSocketServer"
            other_cls.__decorator_params__ = {
                "host": host,
                "port": port
            }
            return other_cls

        return decorator_handler

    @classmethod
    def onopen(cls) -> callable:
        def decorator_handler(func) -> None:
            func.__decorator__ = "onopen"
            cls.__onopen_func = func
            return func

        return decorator_handler

    @classmethod
    def onmessage(cls) -> callable:
        def decorator_handler(func) -> None:
            func.__decorator__ = "onmessage"
            cls.__onmessage_func = func
            return func

        return decorator_handler

    @classmethod
    def onclose(cls) -> callable:
        def decorator_handler(func) -> None:
            func.__decorator__ = "onclose"
            cls.__onclose_func = func
            return func

        return decorator_handler

    @classmethod
    def onerror(cls) -> callable:
        def decorator_handler(func) -> None:
            func.__decorator__ = "onerror"
            cls.__onerror_func = func
            return func

        return decorator_handler


websocket_server = WebSocket.websocket_server
Session = WebSocket.Session
onopen = WebSocket.onopen
onmessage = WebSocket.onmessage
onclose = WebSocket.onclose
onerror = WebSocket.onerror
