"""
inputs/PygameController.py - Pygame Joystick Input Provider with Corrected Axes
"""
import math
from multiprocessing import Queue
import pygame

from inputs.InputProvider import InputProvider, CommandFrame


class PygameController(InputProvider):
    def __init__(self, deadzone: float = 0.2):
        pygame.init()
        pygame.joystick.init()
        
        self.screen = pygame.display.set_mode((400, 300))
        pygame.display.set_caption("JackBot Controller Input")
        self.font = pygame.font.Font(None, 24)
        
        self.deadzone = deadzone
        self.joysticks = {}
        self.clock = pygame.time.Clock()
        self.current_command = CommandFrame(state="idle", vector_dirmov=[0.0, 0.0, 0.0])

    def _normalize_input(self, x: float, y: float) -> tuple[float, float]:
        magnitude = math.sqrt(x**2 + y**2)
        if magnitude > 1.0:
            x /= magnitude
            y /= magnitude
        return x, y

    def update(self) -> CommandFrame:
        vx, vy, omega = 0.0, 0.0, 0.0
        emote = ""

        # Process Pygame events
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pass
            elif event.type == pygame.JOYDEVICEADDED:
                joy = pygame.joystick.Joystick(event.device_index)
                self.joysticks[joy.get_instance_id()] = joy
                print(f"Joystick {joy.get_instance_id()} connected: {joy.get_name()}")
            elif event.type == pygame.JOYDEVICEREMOVED:
                if event.instance_id in self.joysticks:
                    print(f"Joystick {event.instance_id} disconnected")
                    del self.joysticks[event.instance_id]
            elif event.type == pygame.JOYBUTTONDOWN:
                if event.button == 3:
                    emote = "wave"

        # Read active controller stick states
        if self.joysticks:
            for joystick in self.joysticks.values():
                # Axis 1 = Left Stick Up/Down (Negated so Forward = +1.0)
                # Axis 0 = Left Stick Left/Right (Right = +1.0)
                raw_vx = -joystick.get_axis(1)  # Forward (+) / Backward (-)
                raw_vy = joystick.get_axis(0)   # Right (+) / Left (-)
                
                # Check Axis 3 first for Right Stick X, fallback to Axis 2
                num_axes = joystick.get_numaxes()
                raw_omega = joystick.get_axis(3) if num_axes > 3 else joystick.get_axis(2)

                norm_vx, norm_vy = self._normalize_input(raw_vx, raw_vy)

                # Apply deadzones
                if abs(raw_omega) >= self.deadzone:
                    omega = raw_omega
                if abs(norm_vx) >= self.deadzone or abs(norm_vy) >= self.deadzone:
                    vx, vy = norm_vx, norm_vy

                # Axis snapping for clean cardinal walking
                if 0.8 < vx and -0.2 < vy < 0.2:
                    vx, vy = 1.0, 0.0
                elif vx < -0.8 and -0.2 < vy < 0.2:
                    vx, vy = -1.0, 0.0
                elif -0.2 < vx < 0.2 and 0.8 < vy:
                    vx, vy = 0.0, 1.0
                elif -0.2 < vx < 0.2 and vy < -0.8:
                    vx, vy = 0.0, -1.0

        # Determine robot state
        state = "walking" if (vx != 0.0 or vy != 0.0 or omega != 0.0) else "idle"

        self.current_command = CommandFrame(
            state=state,
            vector_dirmov=[vx, vy, omega],
            emote=emote
        )

        # Render visual display
        self.screen.fill((30, 30, 30))
        
        axis_count = list(self.joysticks.values())[0].get_numaxes() if self.joysticks else 0
        
        lines = [
            f"Connected Joysticks: {len(self.joysticks)} (Axes: {axis_count})",
            f"Robot State: {self.current_command.state.upper()}",
            f"Vx (Fwd/Bwd): {vx:>6.2f}",
            f"Vy (Strafe) : {vy:>6.2f}",
            f"Omega (Yaw) : {omega:>6.2f}",
        ]
        for i, line in enumerate(lines):
            txt = self.font.render(line, True, (255, 255, 255))
            self.screen.blit(txt, (20, 20 + i * 30))
        
        pygame.display.flip()
        return self.current_command

    def get_command(self) -> CommandFrame:
        return self.update()

    def stop(self) -> None:
        pygame.quit()


def controller_loop(control_queue: Queue):
    provider = PygameController()
    try:
        while True:
            cmd = provider.get_command()
            control_queue.put({
                "state": cmd.state,
                "dirmov": cmd.vector_dirmov,
                "emote": cmd.emote
            })
            provider.clock.tick(30)
    finally:
        provider.stop()