from __future__ import annotations

import struct
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Optional

import numpy as np
import pygame

from src.vec import Vec

if TYPE_CHECKING:
    from src.car import Car
    from src.game import Game


def iter_unpack(format, data):
    nbr_bytes = struct.calcsize(format)
    return struct.unpack(format, data[:nbr_bytes]), data[nbr_bytes:]


@dataclass
class Snapshot:
    controls: tuple[bool, bool, bool, bool] = (False, False, False, False)
    position: Vec = field(default_factory=Vec)
    direction: Vec = field(default_factory=Vec)
    speed: float = 0
    raycast_distances: list[float] | tuple[float, ...] = field(
        default_factory=list)
    image: Optional[np.ndarray] = None

    def pack(self):
        data: bytes = b""
        data += struct.pack(">BBBB", *self.controls)
        data += struct.pack(
            ">fffff",
            self.position.x,
            self.position.y,
            self.direction.x,
            self.direction.y,
            self.speed,
        )

        nbr_raycasts: int = len(self.raycast_distances)
        data += struct.pack(f">B{nbr_raycasts}f",
                            nbr_raycasts, *self.raycast_distances)

        if self.image is not None:
            data += struct.pack(">II",
                                self.image.shape[0], self.image.shape[1])
            data += self.image.tobytes()
        else:
            data += struct.pack(">II", 0, 0)

        return data

    @staticmethod
    def unpack(data: bytes) -> Snapshot:
        controls, data = iter_unpack(">BBBB", data)
        (x, y, dx, dy, s), data = iter_unpack(">fffff", data)
        position = Vec(x, y)
        direction = Vec(dx, dy)
        speed = s

        (nbr_raycasts,), data = iter_unpack(">B", data)
        raycast_distances, data = iter_unpack(f">{nbr_raycasts}f", data)

        (h, w), data = iter_unpack(">II", data)

        if h * w > 0:
            image = np.frombuffer(data, np.uint8).reshape(h, w, 3)
        else:
            image = None

        return Snapshot(
            controls=controls,
            position=position,
            direction=direction,
            speed=speed,
            raycast_distances=raycast_distances,
            image=image,
        )

    @staticmethod
    def from_car(car: Car) -> Snapshot:
        return Snapshot(
            controls=(
                car.forward,
                car.backward,
                car.left,
                car.right
            ),
            position=car.pos.copy(),
            direction=car.direction.copy(),
            speed=car.speed,
            raycast_distances=car.rays.copy(),
            image=None
        )

    def apply(self, car: Car):
        car.pos = self.position.copy()
        car.direction = self.direction.copy()
        car.speed = 0

    def add_image(self, game: Game):
        self.image = pygame.surfarray.array3d(game.game_surf)