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 geolocation import Geolocation
from vision_chunk import VisionChunk
from visualization import VisualizationManager, SimMode
from yandex_map import YandexMap

from PIL import Image

import os

class Simulator:
    geo: Geolocation
    yandexMap: YandexMap

    # Менеджер визуализации
    viz_manager: VisualizationManager

    def __init__(self, yandexMap: YandexMap = None):
        self.yandexMap = yandexMap

        # Инициализация переменных для отслеживания траектории
        self.geo = Geolocation(0, 0, 1, 0)
        
        # Создаем папку для изображений, если её нет
        os.makedirs('./images', exist_ok=True)

    def rotate_image_like_drone(self, image: Image.Image, angle: float) -> Image.Image:
        """
        Поворачивает картинку как будто съемка ведется с летящего дрона.
        Выделяет концентрический квадрат, поворачивает его и извлекает результат.
        """
        # Получаем размеры изображения
        width, height = image.size
        square_size = min(width, height)
        off_x = (width - square_size) / 2
        off_y = (height - square_size) / 2

        cropped_image = image.crop((off_x, off_y, off_x + square_size, off_y + 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 update_trajectory(self, dx: float, dy: float):
        """
        Обновляет траекторию полета беспилотника
        """
        # Обновляем текущие координаты
        self.geo.x += dx * self.geo.z
        self.geo.y += dy * self.geo.z

    def update_map(self):
        """
        Обновляет карту траектории полета
        """
        self.viz_manager.update_global_map(self.geo.x, self.geo.y, self.mode)

    def handle(self, dangle: float, velocity: float = 50) -> None:
        """ Сдвиг камеры """
        html = self.yandexMap.driver.find_element(By.TAG_NAME, 'html')

        action = ActionChains(self.yandexMap.driver)
        action.move_to_element_with_offset(html, 200, 200)
        action.click_and_hold()

        self.geo.angle += dangle
        # print(f"  [Simulator] angle: {self.angle / math.pi * 180:.1f}°")
        velocity = max(velocity, 10)
        dx = math.cos(math.pi / 2 + self.geo.angle) * velocity / self.geo.z
        dy = math.sin(math.pi / 2 + self.geo.angle) * velocity / self.geo.z
        # print("  [Simulator] dx, dy:", [dx, dy])
        self.update_trajectory(dx, dy)
        action.move_by_offset(-dx, dy)
        action.release()
        action.perform()
        # print(f"  [Simulator] Position: {self.current_x}, {self.current_y}")
    
    def change_zoom(self, direction: str = 'down'):
        """ Изменение масштаба """
        self.yandexMap.scroll(0.5, 0.5, 2, direction == 'down')
        sleep(0.4)
        if direction == 'down':
            self.geo.z /= 2
        else:
            self.geo.z *= 2

    def get_chunk(self) -> VisionChunk:
        png = self.yandexMap.driver.get_screenshot_as_png()
        im = Image.open(BytesIO(png))

        # Применяем поворот как будто съемка с дрона
        rotated_im = self.rotate_image_like_drone(im, -self.geo.angle)
        return VisionChunk(rotated_im)