import pygame
import random
import math
from collections import deque

# --- Simulation Classes ---

class Task:
    __slots__ = ('source', 'destination', 'item', 'amount')
    def __init__(self, source, destination, item, amount):
        self.source = source
        self.destination = destination
        self.item = item
        self.amount = amount

class Storage:
    def __init__(self, name, pos):
        self.name = name
        self.position = pos  # (x, y) position on screen
        self.inventory = {"item": 1000}  # starting inventory
    
    def add_item(self, item, amount):
        self.inventory[item] = self.inventory.get(item, 0) + amount
    
    def remove_item(self, item, amount):
        if self.inventory.get(item, 0) >= amount:
            self.inventory[item] -= amount
            return True
        return False
    
    def request_transfer(self, destination, item, amount, task_queue):
        if self.inventory.get(item, 0) >= amount:
            # Remove items immediately, simulating reservation.
            self.remove_item(item, amount)
            task = Task(self, destination, item, amount)
            task_queue.append(task)

class Drone:
    __slots__ = ('id', 'x', 'y', 'task', 'speed')
    def __init__(self, id, pos):
        self.id = id
        self.x, self.y = pos
        self.task = None
        self.speed = 2.0  # pixels per frame
    
    def assign_task(self, task):
        self.task = task
        # Set drone's position to the source storage position.
        self.x, self.y = task.source.position
    
    def update(self):
        if self.task:
            dest_x, dest_y = self.task.destination.position
            dx = dest_x - self.x
            dy = dest_y - self.y
            dist = math.hypot(dx, dy)
            if dist < self.speed:
                # Arrived: set position exactly to destination and deliver item.
                self.x, self.y = dest_x, dest_y
                self.task.destination.add_item(self.task.item, self.task.amount)
                self.task = None
            else:
                # Move a step toward the destination.
                self.x += self.speed * dx / dist
                self.y += self.speed * dy / dist

# --- Pygame Simulation ---

def main():
    pygame.init()
    screen_width, screen_height = 800, 600
    screen = pygame.display.set_mode((screen_width, screen_height))
    pygame.display.set_caption("Drone Simulation")
    clock = pygame.time.Clock()
    font = pygame.font.SysFont("Arial", 16)

    # Global simulation state
    task_queue = deque()
    storages = []
    drones = []

    # Unique id counters
    drone_id_counter = 0
    storage_id_counter = 0

    # Create some initial storages at random positions
    for _ in range(5):
        pos = (random.randint(50, screen_width - 50), random.randint(50, screen_height - 50))
        storages.append(Storage(f"Storage-{storage_id_counter}", pos))
        storage_id_counter += 1

    # Create some initial drones at random positions
    for _ in range(20):
        pos = (random.randint(0, screen_width), random.randint(0, screen_height))
        drones.append(Drone(drone_id_counter, pos))
        drone_id_counter += 1

    running = True
    while running:
        # --- Event Handling ---
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_t:
                    # Request a random transfer if at least two storages exist.
                    if len(storages) >= 2:
                        src, dest = random.sample(storages, 2)
                        amount = random.randint(1, 10)
                        src.request_transfer(dest, "item", amount, task_queue)
                elif event.key == pygame.K_d:
                    # Add a new drone at a random position.
                    pos = (random.randint(0, screen_width), random.randint(0, screen_height))
                    drones.append(Drone(drone_id_counter, pos))
                    drone_id_counter += 1
                elif event.key == pygame.K_s:
                    # Add a new storage at a random position.
                    pos = (random.randint(50, screen_width - 50), random.randint(50, screen_height - 50))
                    storages.append(Storage(f"Storage-{storage_id_counter}", pos))
                    storage_id_counter += 1
                elif event.key == pygame.K_i:
                    # Add 50 items to each storage.
                    for s in storages:
                        s.add_item("item", 50)

        # --- Update Simulation ---
        for drone in drones:
            if drone.task is None and task_queue:
                # Assign a task from the queue to the idle drone.
                task = task_queue.popleft()
                drone.assign_task(task)
            drone.update()

        # --- Drawing ---
        screen.fill((30, 30, 30))  # dark background

        # Draw storages (blue circles) with name and inventory.
        for s in storages:
            pygame.draw.circle(screen, (0, 100, 255), s.position, 20)
            info = font.render(f"{s.name}: {s.inventory.get('item', 0)}", True, (255, 255, 255))
            screen.blit(info, (s.position[0] - 30, s.position[1] - 40))

        # Draw drones.
        # Idle drones are drawn in green; drones with tasks (in transit) are red.
        for d in drones:
            color = (255, 0, 0) if d.task else (0, 255, 0)
            pygame.draw.circle(screen, color, (int(d.x), int(d.y)), 5)

        # Debug info: show task queue length, number of drones, storages.
        debug_info = font.render(
            f"Task Queue: {len(task_queue)} | Drones: {len(drones)} | Storages: {len(storages)}",
            True, (255, 255, 0))
        screen.blit(debug_info, (10, 10))

        # Display instructions.
        instructions = [
            "Controls:",
            "T - Request Transfer",
            "D - Add Drone",
            "S - Add Storage",
            "I - Add Items to all Storages"
        ]
        for i, line in enumerate(instructions):
            instr_text = font.render(line, True, (200, 200, 200))
            screen.blit(instr_text, (10, 30 + i * 20))

        pygame.display.flip()
        clock.tick(60)

    pygame.quit()

if __name__ == "__main__":
    main()