import struct
import math
import operator
from functools import reduce    
from color import Color

def bit_slice(value, size, count):
    mask = (2 ** size) - 1
    return [(value >> offset) & mask for offset in range(0, size * count, size)]

def bit_merge(values, size):
    offsets = range(0, len(values) * size, size)
    return reduce(operator.__or__, map(operator.lshift, values, offsets)) 

def triple_slice(triplet):
    values = bit_slice(bit_merge(triplet, 8), 3, 8)
    return [values[0:4], values[4:8]]

def triple_merge(rows):
    values = rows[0] + rows[1]
    return bit_slice(bit_merge(values, 3), 8, 3)

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 repr(self.__dict__)

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

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

        self.color0 = Color.from_565(block[0])
        self.color1 = Color.from_565(block[1])
        self.indices = [bit_slice(row, 2, 4) for row in block[2:6]]

    def write(self, file):
        file.write(struct.pack('<2H4B', 
            self.color0.to_565(), self.color1.to_565(), 
            *(bit_merge(row, 2) for row in self.indices)))


class DXT5Block:
    size = 16

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

        self.alpha0 = 1
        self.alpha1 = 1
        self.alpha_indices = [[0] * 4] * 4

        if file:
            self.read(file)

    def __repr__(self):
        return repr(self.__dict__)

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

        self.alpha0 = block[0] / 0xFF
        self.alpha1 = block[1] / 0xFF
        self.alpha_indices = triple_slice(block[2:5]) + triple_slice(block[5:8])

        self.color0 = Color.from_565(block[8])
        self.color1 = Color.from_565(block[9])
        self.indices = [bit_slice(row, 2, 4) for row in block[10:14]]

    def write(self, file):
        file.write(struct.pack('<2B6B2H4B',
            int(self.alpha0 * 0xFF), int(self.alpha1 * 0xFF),
            *triple_merge(self.alpha_indices[0:2]),
            *triple_merge(self.alpha_indices[2:4]),
            self.color0.to_565(), self.color1.to_565(), 
            *(bit_merge(row, 2) for row in self.indices)))

    def to_dxt1(self):
        dxt1 = DXT1Block()
        dxt1.color0 = self.color0
        dxt1.color1 = self.color1
        dxt1.indices = self.indices
        return dxt1

class BlockTexture:
    def __init__(self, block_type, width = 4, height = 4, file = None):
        self.block_type = block_type
        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(self.block_type(file))

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

class DXT1Texture(BlockTexture):
    def __init__(self, width = 4, height = 4, file = None):
        super().__init__(DXT1Block, width, height, file)

class DXT5Texture(BlockTexture):
    def __init__(self, width = 4, height = 4, file = None):
        super().__init__(DXT5Block, width, height, file)

    def to_dxt1(self):
        dxt1 = DXT1Texture(self.width, self.height)
        dxt1.blocks = [block.to_dxt1() for block in self.blocks]
        # print(f'size:{dxt1.width}x{dxt1.height}, {len(dxt1.blocks)} blocks')
        return dxt1