import socket
import threading
import tkinter as tk
from tkinter import ttk
from PIL import Image, ImageTk
import io

clients = {}
client_frames = {}

def handle_client(conn, addr, client_id):
    print(f"[+] New client: {addr} as {client_id}")
    while True:
        try:
            length_data = conn.recv(4)
            if not length_data:
                break
            img_len = int.from_bytes(length_data, byteorder='big')
            img_data = b''
            while len(img_data) < img_len:
                packet = conn.recv(img_len - len(img_data))
                if not packet:
                    break
                img_data += packet
            if client_id in client_frames:
                image = Image.open(io.BytesIO(img_data))
                tk_image = ImageTk.PhotoImage(image)
                canvas = client_frames[client_id]["canvas"]
                canvas.img = tk_image
                canvas.create_image(0, 0, anchor="nw", image=tk_image)
        except:
            break
    print(f"[-] Disconnected: {addr}")
    del clients[client_id]
    del client_frames[client_id]
    conn.close()
    refresh_client_list()

def refresh_client_list():
    for widget in client_list_frame.winfo_children():
        widget.destroy()
    for client_id in clients:
        btn = ttk.Button(client_list_frame, text=client_id, command=lambda c=client_id: open_viewer(c))
        btn.pack(fill="x", padx=5, pady=2)

def open_viewer(client_id):
    if client_id in client_frames:
        return
    top = tk.Toplevel(root)
    top.title(f"Screen: {client_id}")
    canvas = tk.Canvas(top, width=800, height=600)
    canvas.pack()
    client_frames[client_id] = {"window": top, "canvas": canvas}

def accept_clients():
    while True:
        conn, addr = server_socket.accept()
        client_id = f"{addr[0]}:{addr[1]}"
        clients[client_id] = conn
        threading.Thread(target=handle_client, args=(conn, addr, client_id), daemon=True).start()
        refresh_client_list()

# GUI setup
root = tk.Tk()
root.title("Screen Viewer Server")
root.geometry("300x500")

client_list_frame = tk.Frame(root)
client_list_frame.pack(fill="both", expand=True)

# Networking
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_socket.bind(("0.0.0.0", 9999))
server_socket.listen(5)
print("[*] Server listening on port 9999")

threading.Thread(target=accept_clients, daemon=True).start()
root.mainloop()