import pygame
import sys
from colorama import Fore, Style

class Visualizer:
    def __init__(self, positions : list[tuple[float, float, float]], screen_dimensions : tuple[int, int], gcode_size : tuple[float, float]):
        self.positions = positions
        self.screen = None
        self.clock = None
        self.drawn_lines = []
        self.travel_moves = []
        self.screen_dimensions = screen_dimensions
        self.gcode_size = gcode_size
        self._init_pygame()

        self.prev_x, self.prev_y, self.prev_z = self._transform_coordinates(positions[0])
        self.idx = 1

    def _transform_coordinates(self, pos : tuple[float, float, float]) -> tuple[float, float, float]:
        x, y, z = pos
        # Invert y-axis for screen coordinates
        y = -y

        # Transform the coordinates to fit the screen
        screen_x = x/self.gcode_size[0] * self.screen.get_width() 
        screen_y = y/self.gcode_size[1] * self.screen.get_height()
        screen_y += self.screen.get_height()
        return screen_x, screen_y, z

    def _init_pygame(self) -> None:
        pygame.init()
        self.screen = pygame.display.set_mode(self.screen_dimensions)
        self.clock = pygame.time.Clock()

    def _print_status(self) -> None:
        normalized_x = self.prev_x / self.screen.get_width()
        normalized_y = self.prev_y / self.screen.get_height()
        pen_active = self.prev_z < 0

        # Clear the terminal with ANSI code
        print("\033c", end="")

        # X-axis bar
        x_bar = f"{Fore.GREEN}{'#'*int(normalized_x*100)}{Style.RESET_ALL}{' '*int((1.0-normalized_x)*100)}"
        print(f"X: |{x_bar}| {(normalized_x*100):.2f}%")

        # Y-axis bar
        y_bar = f"{Fore.BLUE}{'#'*int(normalized_y*100)}{Style.RESET_ALL}{' '*int((1.0-normalized_y)*100)}"
        print(f"Y: |{y_bar}| {(normalized_y*100):.2f}%")

        # Pen status
        pen_status = f"{Fore.GREEN}Active{Style.RESET_ALL}" if pen_active else f"{Fore.RED}Inactive{Style.RESET_ALL}"
        print(f"Pen: {pen_status}")

    def visualize(self) -> None:
        running = True
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False

            self.screen.fill((0, 0, 0))

            if self.idx < len(self.positions):
                pos = self.positions[self.idx]
                x, y, z = self._transform_coordinates(pos)
                if self.prev_z < 0 and z < 0:
                    self.drawn_lines.append((self.prev_x, self.prev_y, x, y))
                elif z > 0:
                    self.travel_moves.append((self.prev_x, self.prev_y, x, y))
                self.prev_x, self.prev_y, self.prev_z = x, y, z
                # Draw the current position
                pygame.draw.circle(self.screen, (255, 0, 0), (int(x), int(y)), 5)
                self.idx += 1

            # Draw travel moves
            for travel_move in self.travel_moves:
                x1, y1, x2, y2 = travel_move
                pygame.draw.line(self.screen, (0, 0, 255), (int(x1), int(y1)), (int(x2), int(y2)), 2)

            # Draw drawn lines
            for drawn_line in self.drawn_lines:
                x1, y1, x2, y2 = drawn_line
                pygame.draw.line(self.screen, (255, 255, 255), (int(x1), int(y1)), (int(x2), int(y2)), 2)

            self._print_status()

            pygame.display.flip()
            self.clock.tick(10)