"""
2026-06-30 2100Z mitch

Commercial/factory format of microSD uses "dos" disklabel with a single partition of type 0x0c (W95 FAT32 (LBA)).
See: https://www.pjrc.com/tech/8051/ide/fat32.html
"""
import io
import struct
import logging
from datetime import datetime
logger = logging.getLogger(__name__)

BYTES_PER_SECTOR = 512 # TODO: Figure out how to get this from the device, if at all.

class Disk:
    def __init__(self, device):
        self.device = device
        self.mbr = Mbr(self.read(512), self)
        self.partitions = self.mbr.partitions

    def read(self, count):
        return self.device.read(count)

    def seek(self, position):
        return self.device.seek(position)

    def tell(self):
        return self.device.tell()

    @property
    def u16(self):
        return self.unpack('<H', 2)
    
    @property
    def u8(self):
        return self.unpack('B', 1)

    @property
    def u32(self):
        return self.unpack('<I', 4)

    def bytes(self, count):
        return self.read(count)

    def unpack(self, format_, byte_count):
        return struct.unpack(format_, self.read(byte_count))[0]

class Mbr:
    def __init__(self, bytes_, disk):
        self.disk = disk
        self.boostrap_code = bytes_[:446]
        # 4 Partitions; Only remember the ones that aren't type 0.
        partitions = []
        offset = 446
        for n in range(4):
            partition = Partition(bytes_[offset:offset+16], disk=disk)
            if partition.type != 0x0:
                partitions.append(partition)
            offset += 16
        self.partitions = partitions
        self.boot_signature = bytes_[-2:]

    def __str__(self):
        return f'{self.__class__.__name__}(partitions={len(self.partitions)})'

    __repr__ = __str__

class Partition:
    def __init__(self, bytes_, disk):
        self.disk = disk
        p = struct.unpack('<B3sB3sII', bytes_)
        self.flags     = p[0]
        self.first_chs = ChsAddress(p[1])
        self.type      = p[2]
        self.last_chs  = ChsAddress(p[3])
        self.first_lba = p[4]
        self.sectors   = p[5]
        self.start     = self.first_lba * BYTES_PER_SECTOR
        self.size      = self.sectors * BYTES_PER_SECTOR
        self.end       = self.start + self.size

    def get_fs(self):
        self.disk.seek(self.start)
        return Fat32(self.disk, size=self.size)

    def __str__(self):
        return f'{self.__class__.__name__}(type={hex(self.type)}, start={self.start} bytes, size={self.size} bytes)'

    __repr__ = __str__

class ChsAddress:
    def __init__(self, bytes_):
        self.head, cs_bytes = struct.unpack('<B2s', bytes_)
        self.cylinder = (cs_bytes[0] & 0b1100000) << 2 | cs_bytes[1]
        self.sector = cs_bytes[0] & 0b11111

    def __str__(self):
        return f'{self.__class__.__name__}(head={self.head}, cylinder={self.cylinder}, sector={self.sector})'

    __repr__ = __str__

class Bpb:
    def __init__(self, disk):
        self.disk = disk
        # DOS 2.0
        self.bytes_per_sector    = self.disk.u16
        self.sectors_per_cluster = self.disk.u8
        self.reserved_sectors    = self.disk.u16
        self.fat_count           = self.disk.u8
        self.max_root_dirs       = self.disk.u16
        self.total_sectors       = self.disk.u16
        self.media_descriptor    = self.disk.u8
        self.sectors_per_fat     = self.disk.u16
        # DOS 3.31
        self.sectors_per_track   = self.disk.u16
        self.head_count          = self.disk.u16
        self.hidden_sectors      = self.disk.u32
        self.total_sectors_2     = self.disk.u32
        # FAT32 EBPB
        self.sectors_per_fat     = self.disk.u32
        self.drive_flags         = self.disk.u16
        ver_low, ver_high        = self.disk.bytes(2)
        self.version             = f'{ver_high}.{ver_low}'
        self.cluster_root_start  = self.disk.u32
        self.sector_fs_info      = self.disk.u16
        self.sector_boot_backup  = self.disk.u16
        self.reserved            = self.disk.bytes(12)
        self.drive_number        = self.disk.u8
        self.reserved_2          = self.disk.u8
        self.ext_boot_signature  = self.disk.u8
        self.volume_id           = self.disk.u32
        self.volume_label        = self.disk.bytes(11)
        self.fs_type             = self.disk.bytes(8)


class Fat32:
    READ_ONLY    = 0x01
    HIDDEN       = 0x02
    SYSTEM       = 0x04
    VOLUME_LABEL = 0x08
    SUBDIRECTORY = 0x10
    ARCHIVE      = 0x20
    DEVICE       = 0x40
    RESERVED     = 0x80

    def __init__(self, disk, size):
        self.disk = disk
        self.size = size
        self.parse()

    def parse(self):
        self.partition_offset = self.disk.tell()
        self.jump_instruction = self.disk.bytes(3)
        self.oem_name = self.disk.bytes(8)
        self.bpb = Bpb(self.disk)
        self.disk.seek(self.partition_offset + self.bpb.bytes_per_sector)
        self.cluster_size = self.bpb.sectors_per_cluster * self.bpb.bytes_per_sector
        # FAT
        fat_offset = self.partition_offset + (self.bpb.bytes_per_sector * self.bpb.reserved_sectors)
        self.disk.seek(fat_offset)
        fat = []
        for sector in range(self.bpb.sectors_per_fat * self.bpb.bytes_per_sector // 4):
            fat.append(self.disk.u32)
        # Sanity check.  Make sure the low byte of the first FAT entry is the media_descriptor.
        assert fat[0] & 0xFF == self.bpb.media_descriptor
        # The next entry should be the EOC.
        assert self.is_eoc(fat[1])
        self.fat = fat
        # Compute first cluster offet.
        self.first_cluster_offset = fat_offset + (self.bpb.fat_count * self.bpb.sectors_per_fat * self.bpb.bytes_per_sector)

    def list_root(self):
        self.goto_cluster(self.bpb.cluster_root_start)
        long_filename_ucs2 = bytes()
        for n in range(16):
            short_name = self.disk.bytes(8)
            if short_name[0] == 0x00:
                break
            if short_name[0] == 0xE5:
                self.disk.read(24)
                continue
            short_ext  = self.disk.bytes(3)
            attributes = self.disk.u8
            if attributes == 0x0F:
                # VFAT long filename record.
                is_last = short_name[0] >> 6 & 1
                number = short_name[0] & 0b11111
                lfn_part = short_name[1:] + short_ext
                type_ = self.disk.u8
                assert type_ == 0x00
                checksum = self.disk.u8
                lfn_part += self.disk.read(12)
                first_cluster = self.disk.u16
                assert first_cluster == 0x0000
                lfn_part += self.disk.read(4)
                long_filename_ucs2 = lfn_part + long_filename_ucs2
            else:
                # Regular record.
                self.disk.bytes(2) # Junk.
                time_raw = self.disk.u16
                date_raw = self.disk.u16
                if date_raw:
                    creation_datetime = self.parse_datetime(date_raw, time_raw)
                else:
                    creation_datetime = None
                user_id = self.disk.u8
                group_id = self.disk.u8 # 20
                first_cluster_high = self.disk.bytes(2)
                modified_time_raw = self.disk.u16
                modified_date_raw = self.disk.u16
                if modified_date_raw:
                    modified_datetime = self.parse_datetime(modified_date_raw, modified_time_raw)
                else:
                    modified_datetime = None
                first_cluster_low = self.disk.bytes(2)
                first_cluster = struct.unpack('<I', first_cluster_low + first_cluster_high)[0]
                size = self.disk.u32
                if long_filename_ucs2:
                    # NB: This might not be the right codec for UCS-2.
                    # Options are here: https://docs.python.org/3/library/codecs.html#standard-encodings
                    name = long_filename_ucs2.decode('utf_16_le')
                    # Codepoint 0, if it occurs, marks the end of the name.
                    if '\x00' in name:
                        name, _ = name.split('\x00')
                    # Reset the LFN buffer.
                    long_filename_ucs2 = bytes()
                else:
                    name = (short_name + b'.' + short_ext).decode('ascii')
                # Ignore volume labels.
                if attributes & self.VOLUME_LABEL:
                    continue
                yield File(
                    fs            = self,
                    name          = name,
                    first_cluster = first_cluster,
                    size          = size,
                    attributes    = attributes,
                    ctime         = creation_datetime,
                    )
        else:
            assert 0, "Root dir occupies more than one cluster."

    def parse_datetime(self, date_raw, time_raw):
        year   = (date_raw >> 9) + 1980
        month  = date_raw >> 5 & 0b111
        day    = date_raw & 0b11111
        hour   = time_raw >> 11
        minute = time_raw >> 5 & 0b111111
        second = time_raw & 0b1111 * 2
        return datetime(year, month, day, hour, minute, second)

    def is_eoc(self, i):
        # 0x?FFFFFF8 - 0x?FFFFFFF
        return 0xFFF_FFFF >= i & 0xFFF_FFFF >= 0xFFF_FFF8

    def is_free(self, i):
        # 0x?0000000
        return i & 0xFFF_FFFF == 0

    def goto_cluster(self, n, offset=0):
        """ NB: The first cluster is number 2. """
        lba = self.first_cluster_offset + ((n - 2) * self.cluster_size) + offset
        self.disk.seek(lba)

    def goto_sector(self, n):
        self.disk.seek(self.partition_offset + n * self.bpb.bytes_per_sector)

class File:
    def __init__(self, fs, name, first_cluster, size, attributes, ctime):
        self.fs            = fs
        self.name          = name
        self.first_cluster = first_cluster
        self.size          = size
        self.attributes    = attributes
        self.ctime         = ctime
        self.position      = 0

    def read(self, count):
        if self.position + count > self.size:
            count = self.size - self.position
        if count == 0:
            return b''
        # Run through the cluster chain in the FAT until we have the one that includes our position.
        cluster = self.first_cluster
        n = 0
        for n in range(self.position // self.fs.cluster_size):
            cluster = self.fs.fat[cluster]
        offset = self.position % self.fs.cluster_size
        self.fs.goto_cluster(cluster, offset=offset)
        count_left = count

        buffer = io.BytesIO()
        while count_left:
            bytes_left_in_cluster = self.fs.cluster_size - offset
            bytes_to_read = min(count_left, bytes_left_in_cluster)
            buffer.write(self.fs.disk.read(bytes_to_read))
            count_left -= bytes_to_read
            if count_left:
                last_cluster = cluster
                cluster = self.fs.fat[last_cluster]
                logger.debug(f"{count_left} bytes left, but cluster {last_cluster} is done.  Going to cluster {cluster}.")
                self.fs.goto_cluster(cluster)
                offset = 0
        self.position += count
        return buffer.getvalue() 

    def seek(self, position):
        self.position = position

    def __str__(self):
        return f"{self.__class__.__name__}(name={self.name}, size={self.size}, first_cluster={self.first_cluster}, ctime={self.ctime})"

def list_files(device):
    return Disk(device).partitions[0].get_fs().list_root()

def open_file(device, filename):
    file_by_name = {f.name: f for f in list_files(device)}
    try:
        return file_by_name[filename]
    except KeyError:
        raise FileNotFoundError(f"[Errno 2] No such file or directory: '{filename}'")

if __name__ == '__main__':
    logging.basicConfig(style='{', datefmt='%Y-%m-%d %H:%M:%S', format='{asctime}.{msecs:03.0f}Z {levelname} {name}: {message}', level=logging.DEBUG)
    # Tests.
    seq = open_file(open('/dev/sda', 'rb'), 'seq.u8s')

    # 1. Ensure that we can read a file that inhabits multiple clusters.
    for n in range(10_000):
       i = struct.unpack('<I', seq.read(4))[0]
       assert i == n, f"{i} != {n}"

    # 2. Ensure that we can handle a read that spans clusters.
    seq.seek(seq.fs.cluster_size - 4)
    bytes_ = seq.read(8)
    ints = struct.unpack('<II', bytes_)
    assert ints == (0x1fff, 0x2000)

    # 3. Ensure that, if we try to read beyond the end of the file, we don't.
    seq.seek(0)
    assert len(seq.read(seq.size * 2)) == seq.size 

    # 4. Ensure that sequential reads take up where the last one left off.
    seq.seek(0)
    for n in range(7):
        assert struct.unpack('<I', seq.read(4))[0] == n