import struct
import math
from color import Color

def byte_slice(byte):
    return [byte & 0x03, 
           (byte & 0x0C) >> 2, 
           (byte & 0x30) >> 4,
           (byte & 0xC0) >> 6]

def byte_unslice(indices):
    return indices[0] | (indices[1] << 2) | (indices[2] << 4) | (indices[3] << 6)

class DXT1Texture:
    def __init__(self, width = 4, height = 4, file = None):
        self.width = width
        self.height = height
        self.block_width = math.ceil(width / 4)
        self.block_height = math.ceil(height / 4)
        self.block_count = self.block_width * self.block_height
        self.blocks = []

        if file:
            self.read(file)

    def __repr__(self):
        return f'{self.block_count} blocks: {self.blocks}'

    def read(self, file):
        for x in range(self.block_count):
            self.blocks.append(DXT1Block(file))

    def write(self, file):
        for block in self.blocks:
            block.write(file)

class DXT1Block:
    size = 8

    def __init__(self, file = None):
        self.color0 = Color()
        self.color1 = Color()
        self.indices = [[0] * 4] * 4

        if file:
            self.read(file)

    def __repr__(self):
        return f'color0: ({repr(self.color0)}) color1: ({repr(self.color1)}), indices:{self.indices}'

    def __str__(self):
        return f'color0: {str(self.color0)} color1: {str(self.color1)}, indices:{self.indices}'

    def read(self, file):
        block = struct.unpack_from('<HHBBBB', file.read(self.size))

        self.color0 = Color.from_565(block[0])
        self.color1 = Color.from_565(block[1])
        self.indices = list(map(byte_slice, block[2:6]))

    def write(self, file):
        file.write(struct.pack('<HHBBBB', self.color0.to_565(), self.color1.to_565(), *map(byte_unslice, self.indices)))