import socket
import threading
import json
import hashlib
import struct
import time


class Connector:
    def __init__(self, coordinator_host='127.0.0.1', coordinator_port=5000, local_port=2222):
        self.coordinator_host = coordinator_host
        self.coordinator_port = coordinator_port
        self.local_port = local_port
        self.token = None
        self.connections = {}
        self.conn_counter = 0
        self.lock = threading.Lock()

    def connect_to_coordinator(self):
        self.coord_conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.coord_conn.connect((self.coordinator_host, self.coordinator_port))

        # Login
        self._send_json({'action': 'login', 'account': 'admin'})
        response = self._recv_json()

        if response.get('status') == 'salt':
            salt = response['salt']
            password_hash = hashlib.sha256((salt + "admin_password").encode()).hexdigest()
            self._send_json({'action': 'auth', 'hash': password_hash})
            response = self._recv_json()

            if response.get('status') == 'success':
                self.token = response['token']
                print(f"Authenticated. Token: {self.token}")
                return True
        print("Connection to coordinator failed")
        return False

    def request_service(self, service_name):
        punch_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        punch_socket.bind(('0.0.0.0', 0))
        punch_port = punch_socket.getsockname()[1]
        gathering_port = self.coordinator_port+2
        print(f'punch socket send to {(self.coordinator_host, gathering_port)}')
        punch_socket.sendto(json.dumps({'token': self.token}).encode(), (self.coordinator_host, gathering_port))

        self._send_json({
            'action': 'request_service',
            'service_name': service_name,
            'token': self.token
        })
        time.sleep(3)
        response = self._recv_json()

        if response.get('status') == 'success':
            provider_addr = tuple(response['provider_addr'])
            print(f"Connecting to provider at {provider_addr}")

            # 向对方发送打洞包
            for i in range(20):
                punch_socket.sendto(b'pong pong pong pong', provider_addr)
                time.sleep(0.2)

            # Start listening for incoming connections from provider
            listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            listener.bind(('0.0.0.0', punch_port))
            listener.listen(5)
            print(f"Listening on port {punch_port} for provider connections")

            # Start handler thread to accept provider's connection
            threading.Thread(
                target=self.handle_provider_connection,
                args=(listener, service_name),
                daemon=True
            ).start()
            return True
        print("Failed to request service")
        return False

    def handle_provider_connection(self, listener, service_name):
        # Accept connection from provider
        try:
            provider_sock, addr = listener.accept()
            print(f"Accepted provider connection from {addr}")
            # Start heartbeat monitoring
            threading.Thread(
                target=self.monitor_heartbeats,
                args=(provider_sock,),
                daemon=True
            ).start()

            # Start client listener to accept local clients
            client_listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            client_listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            client_listener.bind(('0.0.0.0', self.local_port))
            client_listener.listen(5)
            print(f"Client listener started on port {self.local_port}")

            try:
                while True:
                    client_sock, addr = client_listener.accept()
                    print(f"New client from {addr}")
                    with self.lock:
                        conn_id = self.conn_counter
                        self.conn_counter += 1

                    threading.Thread(
                        target=self.handle_client_connection,
                        args=(client_sock, provider_sock, conn_id),
                        daemon=True
                    ).start()
            finally:
                client_listener.close()
        finally:
            listener.close()

    def handle_client_connection(self, client_sock, provider_sock, conn_id):
        self.connections[conn_id] = client_sock

        try:
            while True:
                data = client_sock.recv(4096)
                if not data:
                    break
                header = struct.pack("!I B", conn_id, len(data))
                provider_sock.sendall(header + data)
        finally:
            client_sock.close()
            with self.lock:
                if conn_id in self.connections:
                    del self.connections[conn_id]

    def monitor_heartbeats(self, sock):
        last_heartbeat = time.time()
        while True:
            try:
                header = sock.recv(5)
                if not header:
                    break

                conn_id, data_len = struct.unpack("!I B", header)
                data = sock.recv(data_len) if data_len > 0 else b''

                # Check if heartbeat
                if conn_id == 0 and data_len == 0:
                    last_heartbeat = time.time()
                    continue

                # Forward data to client
                with self.lock:
                    if conn_id in self.connections:
                        self.connections[conn_id].sendall(data)
            except ConnectionResetError:
                break

            # Check heartbeat timeout
            if time.time() - last_heartbeat > 10:
                print("Heartbeat timeout")
                break

    def start(self, service_name='ssh'):
        if not self.connect_to_coordinator():
            return

        if self.request_service(service_name):
            while True:
                time.sleep(1)

    def _send_json(self, data):
        self.coord_conn.sendall(json.dumps(data).encode())

    def _recv_json(self):
        data = self.coord_conn.recv(4096)
        return json.loads(data.decode()) if data else None


if __name__ == "__main__":
    connector = Connector(coordinator_host='www.awin-x.top',local_port=2222)
    connector.start(service_name='ssh')