from __future__ import annotations

from dataclasses import dataclass
from typing import Sequence

from tea3.constants import TEA3_LUT_A, TEA3_LUT_B, TEA3_SBOX

MASK32 = 0xFFFFFFFF
MASK64 = 0xFFFFFFFFFFFFFFFF


def _ror8(x: int, n: int = 1) -> int:
    x &= 0xFF
    return ((x >> n) | ((x << (8 - n)) & 0xFF)) & 0xFF


def tea3_reorder_state_byte(b_st_byte: int) -> int:
    b_st_byte &= 0xFF
    b_out = 0
    b_out |= (b_st_byte << 6) & 0x40
    b_out |= (b_st_byte << 1) & 0x20
    b_out |= (b_st_byte << 2) & 0x98
    b_out |= (b_st_byte >> 4) & 0x04
    b_out |= (b_st_byte >> 3) & 0x01
    b_out |= (b_st_byte >> 6) & 0x02
    return b_out & 0xFF


def tea3_state_word_to_newbyte(w_st: int, aw_lut: Sequence[int]) -> int:
    w_st &= 0xFFFF
    b_st0 = w_st & 0xFF
    b_st1 = (w_st >> 8) & 0xFF
    b_out = 0

    for i in range(8):
        b_dist = ((b_st0 >> 5) & 0x03) | ((b_st1 >> 3) & 0x0C)
        if aw_lut[i] & (1 << b_dist):
            b_out |= 1 << i
        b_st0 = _ror8(b_st0, 1)
        b_st1 = _ror8(b_st1, 1)

    return b_out & 0xFF


@dataclass
class Tea3:
    key_register: list[int]
    iv: int
    next_skip_rounds: int

    def __init__(self, frame_number: int, key_register: Sequence[int]):
        if len(key_register) != 10:
            raise ValueError("TEA3 key register must be 10 bytes")
        self.key_register = [b & 0xFF for b in key_register]
        self.iv = self.compute_iv(frame_number)
        self.next_skip_rounds = 51

    @staticmethod
    def compute_iv(frame_number: int) -> int:
        frame_number &= MASK32
        dw_xorred = (frame_number ^ 0xC43A7D51) & MASK32
        dw_xorred = ((dw_xorred << 8) | (dw_xorred >> 24)) & MASK32
        qw_iv = ((frame_number << 32) | dw_xorred) & MASK64
        qw_iv = ((qw_iv >> 8) | ((qw_iv & 0xFF) << 56)) & MASK64
        return qw_iv

    def reset(self, frame_number: int, key_register: Sequence[int]) -> None:
        if len(key_register) != 10:
            raise ValueError("TEA3 key register must be 10 bytes")
        self.key_register = [b & 0xFF for b in key_register]
        self.iv = self.compute_iv(frame_number)
        self.next_skip_rounds = 51

    def step(self) -> None:
        # Step 1: key schedule byte
        b_sbox_out = TEA3_SBOX[self.key_register[7] ^ self.key_register[2]] ^ self.key_register[0]

        self.key_register = self.key_register[1:] + [b_sbox_out & 0xFF]

        # Step 2: state-derived bytes
        b_deriv_byte12 = tea3_state_word_to_newbyte((self.iv >> 8) & 0xFFFF, TEA3_LUT_A)
        b_deriv_byte56 = tea3_state_word_to_newbyte((self.iv >> 40) & 0xFFFF, TEA3_LUT_B)
        b_reord_byte4 = tea3_reorder_state_byte((self.iv >> 32) & 0xFF)

        # Step 3: combine
        b_new_byte = (((self.iv >> 56) & 0xFF) ^ b_reord_byte4 ^ b_deriv_byte12 ^ b_sbox_out) & 0xFF
        b_mix_byte = b_deriv_byte56 & 0xFF

        # Step 4: update 64-bit state
        self.iv = (((self.iv << 8) & MASK64) ^ ((b_mix_byte & 0xFF) << 40)) | b_new_byte
        self.iv &= MASK64

    def next_byte(self) -> int:
        for _ in range(self.next_skip_rounds):
            self.step()
        self.next_skip_rounds = 19
        return (self.iv >> 56) & 0xFF

    def keystream(self, length: int) -> list[int]:
        return [self.next_byte() for _ in range(length)]

    def key_register_view(self) -> tuple[int, ...]:
        return tuple(self.key_register)

    def iv_view(self) -> int:
        return self.iv