# -*- coding: utf-8 -*-
import socket
import threading
import sys

PORT_FOR_A = 8080        # Port for device A to connect
PORT_FOR_BROWSER = 8081  # Port for browser to connect

a_device_socket = None
a_device_lock = threading.Lock()

def tunnel_thread(src, dst):
    """Tunnel thread: continuously forward data"""
    try:
        while True:
            data = src.recv(8192)
            if not data:
                break
            dst.sendall(data)
    except Exception as e:
        pass
    finally:
        try:
            src.close()
        except:
            pass
        try:
            dst.close()
        except:
            pass

def handle_browser(conn, addr):
    """Handle browser connection, establish tunnel with device A"""
    global a_device_socket
    
    print("\n[Browser] Connected from:", addr)
    
    with a_device_lock:
        if not a_device_socket:
            conn.send(b"HTTP/1.1 503 Service Unavailable\r\nContent-Type: text/plain\r\n\r\nDevice A not connected")
            conn.close()
            print("[Browser] Disconnected (Device A not connected):", addr)
            return
    
    t1 = threading.Thread(target=tunnel_thread, args=(conn, a_device_socket))
    t2 = threading.Thread(target=tunnel_thread, args=(a_device_socket, conn))
    t1.start()
    t2.start()
    
    print("[System] Tunnel established between browser and device A")

def handle_a_device(conn):
    """Handle device A connection"""
    global a_device_socket
    
    with a_device_lock:
        if a_device_socket:
            conn.send(b"ERROR: Another A device is already connected")
            conn.close()
            return
        a_device_socket = conn
    
    print("\n[Device A] Connected:", conn.getpeername())
    
    try:
        while True:
            data = conn.recv(1)
            if not data:
                break
    except:
        pass
    finally:
        with a_device_lock:
            a_device_socket = None
        conn.close()
        print("[Device A] Disconnected")

def run_browser_server():
    """Start browser listening port"""
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(('0.0.0.0', PORT_FOR_BROWSER))
    server.listen(10)
    
    print("[System] Browser proxy port started, listening on 0.0.0.0:", PORT_FOR_BROWSER)
    
    while True:
        conn, addr = server.accept()
        t = threading.Thread(target=handle_browser, args=(conn, addr))
        t.start()

def run_a_device_server():
    """Start device A connection port"""
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(('0.0.0.0', PORT_FOR_A))
    server.listen(1)
    
    print("[System] Device A port started, listening on 0.0.0.0:", PORT_FOR_A)
    
    while True:
        conn, addr = server.accept()
        t = threading.Thread(target=handle_a_device, args=(conn,))
        t.start()

def main():
    print("=" * 60)
    print("           Proxy Server (Run on Computer B)")
    print("=" * 60)
    print("Device A connection port:", PORT_FOR_A)
    print("Browser access port:", PORT_FOR_BROWSER)
    print("=" * 60)
    print("Waiting for connections...\n")
    
    t1 = threading.Thread(target=run_a_device_server)
    t2 = threading.Thread(target=run_browser_server)
    t1.start()
    t2.start()
    
    t1.join()
    t2.join()

if __name__ == '__main__':
    main()