# References:
# - https://git.proxmox.com/?p=pve-qemu.git;a=blob;f=vma_spec.txt
# - https://lists.gnu.org/archive/html/qemu-devel/2013-02/msg03667.html
from __future__ import annotations

import hashlib
import struct
from collections import defaultdict
from functools import lru_cache
from typing import TYPE_CHECKING, BinaryIO
from uuid import UUID

from dissect.util import ts
from dissect.util.stream import AlignedStream

from dissect.archive.c_vma import VMA_EXTENT_MAGIC, VMA_MAGIC, c_vma
from dissect.archive.exceptions import InvalidHeaderError

if TYPE_CHECKING:
    from collections.abc import Iterator
    from datetime import datetime


class VMA:
    """Proxmox VMA.

    Parse and provide a readable object for devices in a Proxmox VMA backup file.
    VMA is designed to be streamed for extraction, so we need to do some funny stuff to create a readable
    object from it. Performance is not optimal, so it's generally advised to extract a VMA instead.
    The ``vma-extract`` utility can be used for that.
    """

    def __init__(self, fh: BinaryIO):
        self.fh = fh

        fh.seek(0)
        self.header = c_vma.VmaHeader(fh)
        if self.header.magic != VMA_MAGIC:
            raise InvalidHeaderError("Invalid VMA header magic")

        fh.seek(0)
        header_data = bytearray(fh.read(self.header.header_size))
        header_data[32:48] = b"\x00" * 16
        if hashlib.md5(header_data).digest() != self.header.md5sum:
            raise InvalidHeaderError("Invalid VMA checksum")

        self.version = self.header.version
        self.uuid = UUID(bytes=self.header.uuid)

        blob_start = self.header.blob_buffer_offset
        blob_end = self.header.blob_buffer_offset + self.header.blob_buffer_size
        self._blob = memoryview(bytes(header_data))[blob_start:blob_end]

        blob_offset = 1
        self._blob_data: dict[int, bytes] = {}
        while blob_offset + 2 <= self.header.blob_buffer_size:
            # The header is in big endian, but this is little...
            size = struct.unpack("<H", self._blob[blob_offset : blob_offset + 2])[0]
            if blob_offset + 2 + size <= blob_end:
                self._blob_data[blob_offset] = self._blob[blob_offset + 2 : blob_offset + 2 + size].tobytes()
            blob_offset += size + 2

        self._config: dict[str, bytes] = {}
        for conf_name, conf_data in zip(self.header.config_names, self.header.config_data, strict=False):
            if (conf_name, conf_data) == (0, 0):
                continue

            self._config[self.blob_string(conf_name)] = self.blob_data(conf_data)

        self._devices: dict[int, Device] = {}
        for dev_id, dev_info in enumerate(self.header.dev_info):
            if dev_id == 0 or dev_info.devname_ptr == 0:
                continue

            self._devices[dev_id] = Device(self, dev_id, self.blob_string(dev_info.devname_ptr), dev_info.size)

        self._extent = lru_cache(65536)(self._extent)

    @property
    def creation_time(self) -> datetime:
        return ts.from_unix(self.header.ctime)

    def blob_data(self, offset: int) -> bytes:
        if offset not in self._blob_data:
            raise KeyError(f"No blob data for offset {offset}")
        return self._blob_data[offset]

    def blob_string(self, offset: int) -> str:
        return self.blob_data(offset).decode().rstrip("\x00")

    def config(self, name: str) -> bytes:
        return self._config[name]

    def configs(self) -> dict[str, bytes]:
        return self._config

    def device(self, dev_id: int) -> Device:
        return self._devices[dev_id]

    def devices(self) -> list[Device]:
        return list(self._devices.values())

    def _extent(self, offset: int) -> Extent:
        return Extent(self.fh, offset)

    def extents(self) -> Iterator[Extent]:
        offset = self.header.header_size
        while True:
            try:
                extent = self._extent(offset)
            except EOFError:
                break

            yield extent

            offset += c_vma.VMA_EXTENT_HEADER_SIZE + extent.size


class Device:
    def __init__(self, vma: VMA, dev_id: int, name: str, size: int):
        self.vma = vma
        self.id = dev_id
        self.name = name
        self.size = size

    def __repr__(self) -> str:
        return f"<Device id={self.id} name={self.name} size={self.size}>"

    def open(self) -> DeviceDataStream:
        return DeviceDataStream(self)


class Extent:
    def __init__(self, fh: BinaryIO, offset: int):
        self.fh = fh
        self.offset = offset
        self.data_offset = offset + c_vma.VMA_EXTENT_HEADER_SIZE

        self.fh.seek(offset)
        header_data = bytearray(fh.read(c_vma.VMA_EXTENT_HEADER_SIZE))
        self.header = c_vma.VmaExtentHeader(header_data)
        if self.header.magic != VMA_EXTENT_MAGIC:
            raise InvalidHeaderError("Invalid VMA extent header magic")

        header_data[24:40] = b"\x00" * 16
        if hashlib.md5(header_data).digest() != self.header.md5sum:
            raise InvalidHeaderError("Invalid VMA extent checksum")

        self.uuid = UUID(bytes=self.header.uuid)
        self.size = self.header.block_count * c_vma.VMA_BLOCK_SIZE

        # Keep track of the lowest and highest cluster we have for any device
        # We can use this to speed up extent lookup later on
        # There are at most 59 entries, so safe to parse ahead of use
        self._min = {}
        self._max = {}
        self.blocks = defaultdict(list)
        block_offset = self.data_offset
        for block_info in self.header.blockinfo:
            cluster_num = block_info & 0xFFFFFFFF
            dev_id = (block_info >> 32) & 0xFF
            mask = block_info >> (32 + 16)

            if dev_id == 0:
                continue

            if dev_id not in self._min:
                self._min[dev_id] = cluster_num
                self._max[dev_id] = cluster_num
            elif cluster_num < self._min[dev_id]:
                self._min[dev_id] = cluster_num
            elif cluster_num > self._max[dev_id]:
                self._max[dev_id] = cluster_num

            self.blocks[dev_id].append((cluster_num, mask, block_offset))

            if mask == 0xFFFF:
                block_offset += 16 * c_vma.VMA_BLOCK_SIZE
            elif mask == 0:
                pass
            else:
                block_offset += mask.bit_count() * c_vma.VMA_BLOCK_SIZE

    def __repr__(self) -> str:
        return f"<Extent offset=0x{self.offset:x} size=0x{self.size:x}>"


class DeviceDataStream(AlignedStream):
    def __init__(self, device: Device):
        self.device = device
        self.vma = device.vma
        super().__init__(size=device.size, align=c_vma.VMA_CLUSTER_SIZE)

    def _read(self, offset: int, length: int) -> bytes:
        cluster_offset = offset // c_vma.VMA_CLUSTER_SIZE
        cluster_count = (length + c_vma.VMA_CLUSTER_SIZE - 1) // c_vma.VMA_CLUSTER_SIZE
        block_count = (length + c_vma.VMA_BLOCK_SIZE - 1) // c_vma.VMA_BLOCK_SIZE

        result = []
        for _, mask, block_offset in _iter_clusters(self.vma, self.device.id, cluster_offset, cluster_count):
            read_count = min(block_count, 16)

            # Optimize reading fully set and fully sparse masks
            if mask == 0xFFFF:
                self.vma.fh.seek(block_offset)
                result.append(self.vma.fh.read(c_vma.VMA_BLOCK_SIZE * read_count))
            elif mask == 0:
                result.append(b"\x00" * read_count * c_vma.VMA_BLOCK_SIZE)
            else:
                self.vma.fh.seek(block_offset)
                for allocated, count in _iter_mask(mask, read_count):
                    if allocated:
                        result.append(self.vma.fh.read(c_vma.VMA_BLOCK_SIZE * count))
                    else:
                        result.append(b"\x00" * count * c_vma.VMA_BLOCK_SIZE)

            block_count -= read_count
            if block_count == 0:
                break

        return b"".join(result)


def _iter_clusters(vma: VMA, dev_id: int, cluster: int, count: int) -> Iterator[tuple[int, int, int]]:
    # Find clusters and starting offsets in all extents
    temp = {}
    end = cluster + count

    for extent in vma.extents():
        if dev_id not in extent.blocks:
            continue

        if end < extent._min[dev_id] or cluster > extent._max[dev_id]:
            continue

        for cluster_num, mask, block_offset in extent.blocks[dev_id]:
            if cluster_num == cluster:
                yield cluster_num, mask, block_offset
                cluster += 1

                while cluster in temp:
                    yield temp[cluster]
                    del temp[cluster]
                    cluster += 1
            elif cluster < cluster_num <= end:
                temp[cluster_num] = (cluster_num, mask, block_offset)

            if cluster == end:
                break

        if cluster == end:
            break

    while cluster in temp:
        yield temp[cluster]
        del temp[cluster]
        cluster += 1


def _iter_mask(mask: int, length: int) -> Iterator[tuple[int, int]]:
    # Yield consecutive bitmask values
    current_status = mask & 1
    current_count = 0

    for bit_idx in range(length):
        status = (mask & (1 << bit_idx)) >> bit_idx
        if status == current_status:
            current_count += 1
        else:
            yield current_status, current_count
            current_status = status
            current_count = 1

    if current_count:
        yield current_status, current_count
