import pygame
import sys
import time
import asyncio
import websockets

# Initialize Pygame
pygame.init()

# Set up the screen
screen_width = 600
screen_height = 400
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption("Joystick Control")

# Set up colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
GRAY = (150, 150, 150)

async def websocket_thread(websocket, axis_x, axis_y):
    waiting = await websocket.ping()
    latency = await waiting # This extra await gives async the time to send the websocket stuff I think
    # It's a bit hacky, but it works
    # print(f"latency: {latency}")
    if axis_y >= 0:
        await websocket.send(f"U{axis_y}")
        print("packet?")
    elif axis_y < 0:
        await websocket.send(f"D{-axis_y}")
        print("packet?")
    if axis_x >= 0:
        await websocket.send(f"L{axis_x}")
        print("packet?")
    elif axis_x < 0:
        await websocket.send(f"R{-axis_x}")
        print("packet?")


# Main loop
async def main():
    running = True
    # Get joystick
    joystick_count = pygame.joystick.get_count()
    if joystick_count == 0:
        print("No joystick detected")
        pygame.quit()
        sys.exit()
    joystick = pygame.joystick.Joystick(0)
    joystick.init()

    # Set up font
    font = pygame.font.SysFont(None, 36)

    # Set up clock
    clock = pygame.time.Clock()

    # Initial request status and toggle status
    request_status = False
    request_text = "Not Enabled"
    toggle_status = False  # False for disabled, True for enabled

    # IP and port for the request
    ip_address = "10.0.0.29"
    port = "5000"

    # Flag to track focus on IP and Port entry
    ip_focused = False
    port_focused = False
    async with websockets.connect("ws://10.0.0.29:8765") as websocket:
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.MOUSEBUTTONDOWN:
                    # Check if mouse click is inside IP entry box
                    if ip_rect.collidepoint(event.pos):
                        ip_focused = True
                        port_focused = False
                    # Check if mouse click is inside Port entry box
                    elif port_rect.collidepoint(event.pos):
                        port_focused = True
                        ip_focused = False
                    # Check if mouse click is inside toggle button
                    elif toggle_rect.collidepoint(event.pos):
                        toggle_status = not toggle_status
                        request_status = False  # Reset request status when toggling
                    else:
                        ip_focused = False
                        port_focused = False
                elif event.type == pygame.KEYDOWN:
                    # Capture user input for IP and Port
                    if ip_focused:
                        if event.key == pygame.K_BACKSPACE:
                            ip_address = ip_address[:-1]
                        if event.key == pygame.K_TAB:
                            port_focused = True
                            ip_focused = False
                        elif event.unicode.isdigit() or event.unicode == '.':
                            ip_address += event.unicode
                    elif port_focused:
                        if event.key == pygame.K_BACKSPACE:
                            port = port[:-1]
                        if event.key == pygame.K_TAB:
                            port_focused = False
                            ip_focused = True
                        elif event.unicode.isdigit():
                            port += event.unicode

            # Clear the screen
            screen.fill(BLACK)

            axis_x = int(joystick.get_axis(0) * 100)
            axis_y = int(joystick.get_axis(1) * 100)

            if axis_x > -10 and axis_x < 10:
                axis_x = 0
            if axis_y > -10 and axis_y < 10:
                axis_y = 0

            # Draw IP entry box
            ip_text = font.render("IP:", True, WHITE)
            screen.blit(ip_text, (10, 25))
            ip_rect = pygame.Rect(70, 20, 150, 36)
            pygame.draw.rect(screen, GRAY if ip_focused else WHITE, ip_rect, 2)
            ip_text_surface = font.render(ip_address, True, WHITE)
            screen.blit(ip_text_surface, (ip_rect.x + 5, ip_rect.y + 5))

            # Draw Port entry box
            port_text = font.render("Port:", True, WHITE)
            screen.blit(port_text, (10, 75))
            port_rect = pygame.Rect(70, 70, 150, 36)
            pygame.draw.rect(screen, GRAY if port_focused else WHITE, port_rect, 2)
            port_text_surface = font.render(port, True, WHITE)
            screen.blit(port_text_surface, (port_rect.x + 5, port_rect.y + 5))

            # Draw toggle button
            toggle_text = font.render("Enabled" if toggle_status else "Disabled", True, WHITE)
            toggle_rect = pygame.Rect(10, 130, 150, 36)
            pygame.draw.rect(screen, GREEN if toggle_status else RED, toggle_rect)
            pygame.draw.rect(screen, WHITE, toggle_rect, 2)
            screen.blit(toggle_text, (toggle_rect.x + 10, toggle_rect.y + 5))

            # Draw request status
            status_text = font.render("Status:", True, WHITE)
            screen.blit(status_text, (320, 10))
            if request_status:
                pygame.draw.circle(screen, GREEN, (455, 23), 12)
            else:
                pygame.draw.circle(screen, RED, (455, 23), 12)

            if request_text != "":
                request_error_text = font.render(request_text, True, WHITE)
                screen.blit(request_error_text, (350, 55))

            # Draw the ball speed
            speed_text = font.render(f"X Speed: {axis_x}%, Y Speed: {axis_y}%", True, WHITE)
            screen.blit(speed_text, (10, 300))


            # Update the display
            pygame.display.flip()

            # Send data request every 200 ms if toggle is enabled
            if toggle_status and time.time() % 0.01 < 0.01:
                if ip_address and port:
                    await websocket_thread(websocket, axis_x, axis_y)
                else:
                    request_text = "IP and Port empty"
                asyncio.wait(0.01)

            # Cap the frame rate
            clock.tick(60)

    # Quit Pygame
    pygame.quit()

asyncio.run(main())