import asyncio
import websockets
import json

# Guacamole protocol message formatting functions
def encode_guac_message(opcode, *params):
    parts = [f"{len(p)}.{p}" for p in params]
    return f"{opcode},{','.join(parts)};"

def decode_guac_message(message):
    parts = message.strip(';').split(',')
    opcode = parts[0]
    params = []
    
    current_pos = 1
    while current_pos < len(parts):
        length_str = parts[current_pos].split('.')[0]
        length = int(length_str)
        # Note: A real implementation would handle string content more carefully
        param = parts[current_pos][len(length_str) + 1:]
        params.append(param)
        current_pos += 1
    
    return opcode, params

# Simplified VM and User classes for state management
class VM:
    def __init__(self, vm_id, display_name):
        self.vm_id = vm_id
        self.display_name = display_name
        self.users = {} # {username: WebSocket connection}
        self.turn_queue = []
        self.chat_history = []
        self.screen_size = (800, 600)

class User:
    def __init__(self, username, rank=0):
        self.username = username
        self.rank = rank

# Server state
vms = {
    "vm-001": VM("vm-001", "Linux Desktop"),
    "vm-002": VM("vm-002", "Windows 98")
}

# Main server logic
async def handle_client(websocket, path):
    # Enforce guacamole subprotocol
    if "guacamole" not in websocket.subprotocol:
        await websocket.close()
        print("Connection rejected: Incorrect subprotocol.")
        return

    # Handshake phase
    print("New connection established. Starting handshake...")
    user = User(username=None)
    vm_id = None
    
    # Send NOPs periodically (this is a simplified example)
    async def nop_keepalive():
        while True:
            await asyncio.sleep(10) # 10-second interval
            if websocket.open:
                await websocket.send(encode_guac_message("nop"))
            else:
                break
    
    asyncio.ensure_future(nop_keepalive())

    try:
        async for message in websocket:
            opcode, params = decode_guac_message(message)
            print(f"Received opcode: {opcode} with params: {params}")

            # Handshake opcodes
            if opcode == "list":
                list_message = "list,"
                for v in vms.values():
                    # For a real server, thumbnail would be base64-encoded image data
                    list_message += f"{len(v.vm_id)}.{v.vm_id},"
                    list_message += f"{len(v.display_name)}.{v.display_name},"
                    list_message += f"{4}.THMB;" # Mock thumbnail
                await websocket.send(list_message)
            
            elif opcode == "rename":
                requested_name = params[0] if params else "guest"
                user.username = requested_name # Simplified, server would check for availability
                
                # Send back a Client Renamed event
                response = encode_guac_message("rename", "0", "0", user.username)
                await websocket.send(response)

            elif opcode == "connect":
                requested_vm_id = params[0]
                if requested_vm_id in vms:
                    vm = vms[requested_vm_id]
                    vm_id = requested_vm_id
                    vm.users[user.username] = websocket
                    print(f"User {user.username} connected to VM {vm_id}")
                    
                    # Connection success message
                    await websocket.send(encode_guac_message("connect", "1"))
                    
                    # Announce new user to others in the VM
                    for other_user_ws in vm.users.values():
                        if other_user_ws != websocket:
                            await other_user_ws.send(encode_guac_message("adduser", user.username, "0"))
                    
                    # Announce VM screen size
                    width, height = vm.screen_size
                    await websocket.send(encode_guac_message("size", "0", str(width), str(height)))
                    
                    # Send mock framebuffer update (png)
                    # A real implementation would send a base64 encoded image
                    await websocket.send(encode_guac_message("png", "0", "0", "0", "0", "FAKE_BASE64_IMAGE_DATA"))
                    
                else:
                    await websocket.send(encode_guac_message("connect", "0"))

            # VM interaction opcodes (only after connection)
            elif vm_id:
                vm = vms[vm_id]
                if opcode == "chat":
                    message = params[0]
                    chat_msg = f"{user.username}: {message}"
                    vm.chat_history.append(chat_msg)
                    for other_user_ws in vm.users.values():
                        await other_user_ws.send(encode_guac_message("chat", user.username, message))
                
                # Add more opcode handlers here (e.g., mouse, key, turn)
            
            # Client has sent NOP in response to server
            elif opcode == "nop":
                # Do nothing, NOP is a keepalive
                pass

    except websockets.exceptions.ConnectionClosed as e:
        print(f"Connection closed with code {e.code}: {e.reason}")
    finally:
        # Cleanup on disconnect
        if user.username and vm_id:
            vm = vms[vm_id]
            if user.username in vm.users:
                del vm.users[user.username]
                print(f"User {user.username} disconnected from VM {vm_id}")
                # Announce removed user
                for other_user_ws in vm.users.values():
                    await other_user_ws.send(encode_guac_message("remuser", "1", user.username))

# Main function to start the server
async def main():
    server = await websockets.serve(
        handle_client,
        "0.0.0.0", 8765,
        subprotocols=["guacamole"]
    )
    print("CollabVM server started on ws://localhost:8765")
    await server.wait_closed()

if __name__ == "__main__":
    asyncio.run(main())