import io
import os
from dataclasses import dataclass, field

from .binary import BinaryReader, BinaryWriter
from .typing import mat33, mat33_new, vec3, vec3_new


@dataclass
class DeltaKey:
    delta: float = 0

    @staticmethod
    def size() -> int:
        return 4

    def read(self, read: BinaryReader) -> None:
        self.delta = read.f()

    def write(self, write: BinaryWriter) -> None:
        write.f(self.delta)


@dataclass
class TimeKey:
    time: float = 0

    @staticmethod
    def size() -> int:
        return 4

    def read(self, read: BinaryReader) -> None:
        self.time = read.f()

    def write(self, write: BinaryWriter) -> None:
        write.f(self.time)


@dataclass
class TransformKey:
    position: vec3 = field(default_factory=vec3_new)
    rotation_x: vec3 = field(default_factory=vec3_new)
    rotation_y: vec3 = field(default_factory=vec3_new)

    @staticmethod
    def size() -> int:
        return 36

    def read(self, read: BinaryReader) -> None:
        self.position = read.vec3()
        self.rotation_x = read.vec3()
        self.rotation_y = read.vec3()

    def write(self, write: BinaryWriter) -> None:
        write.vec3(self.position)
        write.vec3(self.rotation_x)
        write.vec3(self.rotation_y)


@dataclass
class WheelKey:
    rotation: float = 0
    steer: float = 0
    z: list[float] = field(default_factory=lambda: [0] * 4)

    @staticmethod
    def size() -> int:
        return 24

    def read(self, read: BinaryReader) -> None:
        self.rotation = read.f()
        self.steer = read.f()
        self.z = [read.f() for _ in range(4)]

    def write(self, write: BinaryWriter) -> None:
        write.f(self.rotation)
        write.f(self.steer)
        [write.f(z) for z in self.z]


@dataclass
class TachoKey:
    vel: vec3 = field(default_factory=vec3_new)
    rpm: int = 0
    gear: int = 0

    @staticmethod
    def size() -> int:
        return 20

    def read(self, read: BinaryReader) -> None:
        self.vel = read.vec3()
        self.rpm = read.u32()
        self.gear = read.u32()

    def write(self, write: BinaryWriter) -> None:
        write.vec3(self.vel)
        write.u32(self.rpm)
        write.u32(self.gear)


@dataclass
class AttackKey:
    attack_cb: int = 0
    attack: int = 0
    attack_car: int = 0
    trail: int = 0
    section: int = 0

    @staticmethod
    def size() -> int:
        return 20

    def read(self, read: BinaryReader) -> None:
        self.attack_cb = read.i32()
        self.attack = read.i32()
        self.attack_car = read.i32()
        self.trail = read.i32()
        self.section = read.i32()

    def write(self, write: BinaryWriter) -> None:
        write.i32(self.attack_cb)
        write.i32(self.attack)
        write.i32(self.attack_car)
        write.i32(self.trail)
        write.i32(self.section)


@dataclass
class VelocityKey:
    vel: vec3 = field(default_factory=vec3_new)

    @staticmethod
    def size() -> int:
        return 12

    def read(self, read: BinaryReader) -> None:
        self.vel = read.vec3()

    def write(self, write: BinaryWriter) -> None:
        write.vec3(self.vel)


@dataclass
class EffectKeySound:
    id: int = -1
    pitch: float = 0
    volume: int = 0
    pan: int = 0
    pan_y: int = 0
    unused: int = 0

    def read(self, read: BinaryReader) -> None:
        self.id = read.i16()
        read.u16()  # pad
        self.pitch = read.f()
        self.volume = read.u8()
        self.pan = read.u8()
        self.pan_y = read.u8()
        self.unused = read.u8()

    def write(self, write: BinaryWriter) -> None:
        write.i16(self.id)
        write.u16(0)  # pad
        write.f(self.pitch)
        write.u8(self.volume)
        write.u8(self.pan)
        write.u8(self.pan_y)
        write.u8(self.unused)


@dataclass
class EffectKeySkid:
    flags: int = 0
    normal: vec3 = field(default_factory=vec3_new)
    position: vec3 = field(default_factory=vec3_new)
    unknown: vec3 = field(default_factory=vec3_new)

    def read(self, read: BinaryReader) -> None:
        self.flags = read.u32()
        self.normal = read.vec3()
        self.position = read.vec3()
        self.unknown = read.vec3()

    def write(self, write: BinaryWriter) -> None:
        write.u32(self.flags)
        write.vec3(self.normal)
        write.vec3(self.position)
        write.vec3(self.unknown)


@dataclass
class EffectKey:
    sounds: list[EffectKeySound] = field(default_factory=lambda: [EffectKeySound()] * 8)
    skids: list[EffectKeySkid] = field(default_factory=lambda: [EffectKeySkid()] * 4)
    rotation: mat33 = field(default_factory=mat33_new)

    @staticmethod
    def size() -> int:
        return 292

    def read(self, read: BinaryReader) -> None:
        self.sounds = [read.any(EffectKeySound) for _ in range(8)]
        self.skids = [read.any(EffectKeySkid) for _ in range(4)]
        self.rotation = read.mat33()

    def write(self, write: BinaryWriter) -> None:
        [write.any(e) for e in self.sounds]
        [write.any(p) for p in self.skids]
        write.mat33(self.rotation)


Key = DeltaKey | TimeKey | TransformKey | WheelKey | TachoKey | AttackKey | VelocityKey | EffectKey
_key_types = (DeltaKey, TimeKey, TransformKey, WheelKey, TachoKey, AttackKey, VelocityKey, EffectKey)


@dataclass
class Channel:
    key_type: type[Key] = DeltaKey
    player_index: int = -1

    def read(self, read: BinaryReader) -> None:
        self.key_type = _key_types[read.u32()]
        param_size = read.u32()
        if param_size == 4:
            self.player_index = read.i32()

    def write(self, write: BinaryWriter) -> None:
        write.u32(_key_types.index(self.key_type))
        if self.player_index == -1:
            write.u32(0)  # param_size
        else:
            write.u32(4)  # param_size
            write.i32(self.player_index)


@dataclass
class Times:
    total: float = 0
    laps: list[float] = field(default_factory=lambda: [0] * 4)
    parts: list[list[float]] = field(default_factory=lambda: [[0] * 3] * 4)

    def read(self, read: BinaryReader) -> None:
        assert read.u32() == 3, "bad times version"
        self.total = read.f()
        self.laps = [read.f() for _ in range(4)]
        self.parts = [[read.f() for _ in range(4)] for _ in range(3)]

    def write(self, write: BinaryWriter) -> None:
        write.u32(3)  # version
        write.f(self.total)
        [write.f(x) for x in self.laps]
        [[write.f(y) for y in x] for x in self.parts]


@dataclass
class Player:
    index: int = 0
    car_name: str = ""
    name: str = ""
    grip: int = 0
    speed: int = 0
    handling: int = 0
    brakes: int = 0
    accel: int = 0

    def read(self, read: BinaryReader, channels: list[Channel]) -> None:
        self.index = read.i32()
        # Read car name.
        assert read.u32() == 32, "bad car name size"
        self.car_name = read.strbytes(32)
        # Read data.
        assert read.u32() == 64, "bad player data size"
        self.name = read.strbytes(50)
        _total_enc = (read.u32() ^ 0x4ADC6BA8) - 0x10000000  # cipher differs in older versions
        _total = read.u32() ^ 0x1DF4E5BB  # cipher differs in older versions
        self.grip = read.u8() ^ 0x7C
        self.speed = read.u8() ^ 0xBD
        self.handling = read.u8() ^ 0x3A
        self.brakes = read.u8() ^ 0xE0
        self.accel = read.u8() ^ 0x1E
        read.u8()  # gap
        # Read channels.
        player_channel_count = read.u32()
        _player_channel_types = [_key_types[read.u32()] for _ in range(player_channel_count)]
        _player_channel_indices = [read.i32() for _ in range(player_channel_count)]

    def write(self, write: BinaryWriter, channels: list[Channel]) -> None:
        write.i32(self.index)
        # Write car name.
        write.u32(32)  # car_name_size
        write.strbytes(self.car_name, 32)
        # Write data.
        write.u32(64)  # data_size
        write.strbytes(self.name, 50)
        total = self.grip + self.speed + self.handling + self.brakes + self.accel
        grip_enc = self.grip ^ 0x7C
        speed_enc = self.speed ^ 0xBD
        handling_enc = self.handling ^ 0x3A
        brakes_enc = self.brakes ^ 0xE0
        accel_enc = self.accel ^ 0x1E
        total_enc = grip_enc + speed_enc + handling_enc + brakes_enc + accel_enc
        write.u32((total_enc + 0x10000000) ^ 0x4ADC6BA8)
        write.u32(total ^ 0x1DF4E5BB)
        write.u8(grip_enc)
        write.u8(speed_enc)
        write.u8(handling_enc)
        write.u8(brakes_enc)
        write.u8(accel_enc)
        write.u8(0)  # gap
        # Write channels.
        player_channels = [(i, c) for i, c in enumerate(channels) if c.player_index == self.index]
        write.u32(len(player_channels))
        for _, c in player_channels:
            write.u32(_key_types.index(c.key_type))
        for c, _ in player_channels:
            write.u32(c)


@dataclass
class Info:
    circuit_index: int = 0
    duration: float = 0
    times: Times | None = None
    players: list[Player] = field(default_factory=list[Player])

    def read(self, read: BinaryReader, channels: list[Channel]) -> None:
        assert read.u32() == 3, "bad info version"
        self.circuit_index = read.u32()
        self.duration = read.f()
        _unknown = read.i32()  # -1
        _delta_channel = read.i32()
        _time_channel = read.i32()
        skip_global_channel_count = read.u32()  # for calculating offsets
        skip_player_count = read.u32()  # for calculating offsets
        # Read times.
        times_size = read.u32()
        if times_size == 72:
            self.times = read.any(Times)
        # Read global channels.
        global_channel_count = read.u32()
        assert skip_global_channel_count == global_channel_count, "global channel count mismatch"
        _global_channel_types = [_key_types[read.u32()] for _ in range(global_channel_count)]
        _global_channel_indices = [read.i32() for _ in range(global_channel_count)]
        # Read players.
        player_count = read.u32()
        assert skip_player_count == player_count, "player count mismatch"
        self.players = [read.any(Player, channels) for _ in range(player_count)]

    def write(self, write: BinaryWriter, channels: list[Channel]) -> None:
        write.i32(3)  # version
        write.i32(self.circuit_index)
        write.f(self.duration)
        write.i32(-1)  # unknown
        write.i32(next((i for i, c in enumerate(channels) if c.key_type == DeltaKey and c.player_index == -1), -1))
        write.i32(next((i for i, c in enumerate(channels) if c.key_type == TimeKey and c.player_index == -1), -1))
        write.u32(sum(1 for c in channels if c.player_index == -1))  # skip_global_channel_count
        write.u32(len(self.players))  # skip_player_count
        # Write times.
        if self.times:
            write.u32(72)  # times_size
            write.any(self.times)
        else:
            write.u32(0)
        # Write global channels.
        global_channels = [(i, c) for i, c in enumerate(channels) if c.player_index == -1]
        write.u32(len(global_channels))
        for _, ch in global_channels:
            write.u32(_key_types.index(ch.key_type))
        for c, _ in global_channels:
            write.u32(c)
        # Write players.
        write.u32(len(self.players))
        [write.any(p, channels) for p in self.players]


@dataclass
class Frame:
    keys: list[Key] = field(default_factory=list[Key])

    def read(self, read: BinaryReader, channels: list[Channel]) -> None:
        for channel in channels:
            key = read.any(channel.key_type)
            self.keys.append(key)

    def write(self, write: BinaryWriter, channels: list[Channel]) -> None:
        for key, channel in zip(self.keys, channels):
            assert type(key) == channel.key_type, "channel key type mismatch"
            write.any(key)


@dataclass
class Sequence:
    channels: list[Channel] = field(default_factory=list[Channel])
    info: Info = field(default_factory=Info)
    frames: list[Frame] = field(default_factory=list[Frame])

    def read(self, read: BinaryReader) -> None:
        assert read.u32() == 0x12345678, "bad sequence first magic value"
        assert read.u32() == 0x78451236, "bad sequence second magic value"
        assert read.u32() == 2, "bad sequence version"
        frames_offset = read.u32()
        info_offset = read.u32()
        channel_count = read.u32()
        frame_count = read.i32()  # -1 in files to calculate from channels
        # Read channels.
        self.channels = [read.any(Channel) for _ in range(channel_count)]
        # Read info.
        assert read.file.tell() == info_offset, "bad sequence info offset"
        info_size = read.u32()
        self.info = read.any(Info, self.channels)
        assert read.file.tell() == info_offset + 4 + info_size, "bad sequence info size"
        # Read frames.
        assert read.file.tell() == frames_offset, "bad sequence frame offset"
        if frame_count == -1:
            frame_count = (read.file.seek(0, os.SEEK_END) - frames_offset) // self._calc_frame_size()
            read.file.seek(frames_offset)
        self.frames = [read.any(Frame, self.channels) for _ in range(frame_count)]

    def write(self, write: BinaryWriter) -> None:
        write.u32(0x12345678)  # magic 1
        write.u32(0x78451236)  # magic 2
        write.u32(2)  # version
        offsets_pos = write.file.tell()
        write.u32(0)  # frames_offset
        write.u32(0)  # info_offset
        write.u32(len(self.channels))
        write.i32(-1)  # frame_count, -1 in files to calculate from channels
        # Write channels.
        [write.any(c) for c in self.channels]
        # Write info.
        info_offset = write.file.tell()
        write.u32(self._calc_info_size())
        write.any(self.info, self.channels)
        # Write frames.
        frames_offset = write.file.tell()
        [write.any(f, self.channels) for f in self.frames]
        # Write header offsets.
        write.file.seek(offsets_pos)
        write.u32(frames_offset)
        write.u32(info_offset)
        write.file.seek(0, os.SEEK_END)

    def _calc_info_size(self) -> int:
        size = 32
        # Add times.
        size += 4
        if self.info.times:
            size += 72
        # Add global channels.
        global_channel_count = sum(1 for c in self.channels if c.player_index == -1)
        size += 4 + global_channel_count * 8
        # Add players.
        size += 4
        for player in self.info.players:
            size += 108
            player_channel_count = sum(1 for c in self.channels if c.player_index == player.index)
            size += 4 + player_channel_count * 8
        return size

    def _calc_frame_size(self) -> int:
        return sum(channel.key_type.size() for channel in self.channels)


@dataclass
class Ghost:
    circuit_index: int = 0
    sequences: list[Sequence] = field(default_factory=list[Sequence])

    def read(self, read: BinaryReader) -> None:
        # Read header.
        assert read.u32() == 1, "bad ghost version"
        assert read.u32() == 0x456F1E9C, "bad ghost first magic value"
        assert read.u32() == 0xFA48EFB4, "bad ghost second magic value"
        sequence_count = read.u32()
        self.circuit_index = read.u32()
        # Read sequences.
        for _ in range(sequence_count):
            total_time = read.f()  # sequence.info.times.total
            buffer = read.bytes(read.u32())
            with BinaryReader(io.BytesIO(buffer), read.encoding) as sequence_read:
                s = sequence_read.any(Sequence)
                assert total_time == s.info.times.total if s.info.times else 0, "sequence total time mismatch"
                self.sequences.append(s)

    def write(self, write: BinaryWriter) -> None:
        # Write header.
        write.u32(1)
        write.u32(0x456F1E9C)
        write.u32(0xFA48EFB4)
        write.u32(len(self.sequences))
        write.u32(self.circuit_index)
        # Write sequences.
        for s in self.sequences:
            write.f(s.info.times.total if s.info.times else 0)
            with BinaryWriter(io.BytesIO(), write.encoding) as sequence_write:
                sequence_write.any(s)
                write.u32(sequence_write.file.tell())
                sequence_write.file.seek(0)
                write.bytes(sequence_write.file.read())
