import math
from io import BytesIO
from time import sleep

from PIL import Image
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.action_chains import ActionChains

from autopilot import Pilot

from enum import Enum

class SimMode(Enum):
    OPERATOR = 1
    AUTONOME = 2

class Simulator:
    driver: webdriver.Chrome
    mode: SimMode

    operatorPilot: Pilot
    autonomePilot: Pilot

    angle: float

    def __init__(self, operatorPilot: Pilot, autonomePilot: Pilot):
        self.mode = SimMode.OPERATOR
        self.operatorPilot = operatorPilot
        self.autonomePilot = autonomePilot

        options = webdriver.ChromeOptions()
        # options.add_experimental_option("detach", True)
        self.driver = webdriver.Chrome(options)
        self.driver.get("https://yandex.ru/maps/43/kazan/?ll=49.103814%2C55.794258&z=14")
        sleep(2)

        action = ActionChains(self.driver)

        # Закрытие левой панели
        action.click(self.driver.find_element(By.CLASS_NAME, 'sidebar-toggle-button'))
        action.move_to_element_with_offset(self.driver.find_element(By.XPATH, "//div[@class='rounded-controls']/div[@class='rounded-controls__child'][5]//button"), 5, 5)
        action.perform()

        # Режим спутника
        action.click(self.driver.find_element(By.CLASS_NAME, '_key_satellite'))
        action.perform()

        self.angle = 0

    def rotate_image_like_drone(self, image: Image.Image, angle: float) -> Image.Image:
        """
        Поворачивает картинку как будто съемка ведется с летящего дрона.
        Выделяет концентрический квадрат, поворачивает его и извлекает результат.
        """
        # Получаем размеры изображения
        width, height = image.size
        square_size = min(width, height)
        cropped_image = image.crop((0, 0, square_size, square_size))
        cropped_image = cropped_image.rotate(angle / math.pi * 180, expand=True)
        
        # Определяем размер концентрического квадрата (80% от минимальной стороны)
        local_square_size = int(square_size / 2 ** 0.5)
        
        # Вычисляем координаты для центрирования квадрата
        left = (cropped_image.width - local_square_size) // 2
        top = (cropped_image.height - local_square_size) // 2
        right = left + local_square_size
        bottom = top + local_square_size
        
        # Вырезаем концентрический квадрат
        final_image = cropped_image.crop((left, top, right, bottom))

        return final_image

    def loop(self):

        html = self.driver.find_element(By.TAG_NAME, 'html')
        action = ActionChains(self.driver)
        velocity = 10

        for i in range(100):
            # Сдвиг камеры
            action = ActionChains(self.driver)
            action.move_to_element_with_offset(html, 200, 200)
            action.click_and_hold()
            dangle = self.operatorPilot.act()
            self.angle += dangle
            dx = math.cos(self.angle) * velocity
            dy = math.sin(self.angle) * velocity
            action.move_by_offset(-dx, dy)
            action.release()
            action.perform()

            # Загрузка скриншота
            png = self.driver.get_screenshot_as_png()
            im = Image.open(BytesIO(png))
            im = im.crop([0, 80, im.width-80, im.height-60])
            
            # Применяем поворот как будто съемка с дрона
            rotated_im = self.rotate_image_like_drone(im, self.angle + math.pi / 2)
            
            # Передаем изображение в AutoPilot для анализа
            self.autonomePilot.handle(rotated_im)
            
            rotated_im.save(f"./images/{i}.png")
            sleep(0.25)

        print("last position: ", self.driver.current_url)