"""
A Pillow loader for .vtf files (aka Valve Texture Format)
REDxEYE <med45c@gmail.com>

Documentation:
  https://developer.valvesoftware.com/wiki/Valve_Texture_Format

The contents of this file are hereby released in the public domain (CC0)
Full text of the CC0 license:
  https://creativecommons.org/publicdomain/zero/1.0/
"""

from __future__ import annotations

import struct
from enum import IntEnum, IntFlag
from math import ceil, log
from typing import IO, NamedTuple

from . import Image, ImageFile


class VTFException(Exception):
    pass


class CompiledVtfFlags(IntFlag):
    # Flags from the *.txt config file
    POINTSAMPLE = 0x00000001
    TRILINEAR = 0x00000002
    CLAMPS = 0x00000004
    CLAMPT = 0x00000008
    ANISOTROPIC = 0x00000010
    HINT_DXT5 = 0x00000020
    PWL_CORRECTED = 0x00000040
    NORMAL = 0x00000080
    NOMIP = 0x00000100
    NOLOD = 0x00000200
    ALL_MIPS = 0x00000400
    PROCEDURAL = 0x00000800

    # These are automatically generated by vtex from the texture data.
    ONEBITALPHA = 0x00001000
    EIGHTBITALPHA = 0x00002000

    # Newer flags from the *.txt config file
    ENVMAP = 0x00004000
    RENDERTARGET = 0x00008000
    DEPTHRENDERTARGET = 0x00010000
    NODEBUGOVERRIDE = 0x00020000
    SINGLECOPY = 0x00040000
    PRE_SRGB = 0x00080000

    UNUSED_00100000 = 0x00100000
    UNUSED_00200000 = 0x00200000
    UNUSED_00400000 = 0x00400000

    NODEPTHBUFFER = 0x00800000

    UNUSED_01000000 = 0x01000000

    CLAMPU = 0x02000000
    VERTEXTEXTURE = 0x04000000
    SSBUMP = 0x08000000

    UNUSED_10000000 = 0x10000000

    BORDER = 0x20000000

    UNUSED_40000000 = 0x40000000
    UNUSED_80000000 = 0x80000000


class VtfPF(IntEnum):
    # NONE = -1
    RGBA8888 = 0
    # ABGR8888 = 1
    RGB888 = 2
    BGR888 = 3
    # RGB565 = 4
    I8 = 5
    IA88 = 6
    # P8 = 7
    A8 = 8
    # RGB888_BLUESCREEN = 9
    # BGR888_BLUESCREEN = 10
    # ARGB8888 = 11
    BGRA8888 = 12
    DXT1 = 13
    DXT3 = 14
    DXT5 = 15
    # BGRX8888 = 16
    # BGR565 = 17
    # BGRX5551 = 18
    # BGRA4444 = 19
    DXT1_ONEBITALPHA = 20
    # BGRA5551 = 21
    UV88 = 22
    # UVWQ8888 = 23
    # RGBA16161616F = 24
    # RGBA16161616 = 25
    # UVLX8888 = 26


class VTFHeader(NamedTuple):
    header_size: int
    width: int
    height: int
    flags: int
    frames: int
    first_frames: int
    reflectivity_r: float
    reflectivity_g: float
    reflectivity_b: float
    bumpmap_scale: float
    pixel_format: int
    mipmap_count: int
    low_pixel_format: int
    low_width: int
    low_height: int
    depth: int = 0
    resource_count: int = 0


BLOCK_COMPRESSED = (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA, VtfPF.DXT3, VtfPF.DXT5)
HEADER_V70 = "<I2HI2H4x3f4xfIbI2b"
HEADER_V72 = "<I2HI2H4x3f4xfIbI2bH"
HEADER_V73 = "<I2HI2H4x3f4xfIbI2bH3xI8x"


def _get_texture_size(pixel_format: int, width: int, height: int) -> int:
    if pixel_format in (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA):
        return width * height // 2
    elif pixel_format in (
        VtfPF.DXT3,
        VtfPF.DXT5,
        VtfPF.A8,
        VtfPF.I8,
    ):
        return width * height
    elif pixel_format in (VtfPF.UV88, VtfPF.IA88):
        return width * height * 2
    elif pixel_format in (VtfPF.RGB888, VtfPF.BGR888):
        return width * height * 3
    elif pixel_format == VtfPF.RGBA8888:
        return width * height * 4
    msg = f"Unsupported VTF pixel format: {pixel_format}"
    raise VTFException(msg)


def _get_mipmap_count(width: int, height: int) -> int:
    mip_count = 1
    while True:
        mip_width = width >> mip_count
        mip_height = height >> mip_count
        if mip_width == 0 and mip_height == 0:
            return mip_count
        mip_count += 1


def _write_image(fp: IO[bytes], im: Image.Image, pixel_format: VtfPF) -> None:
    args: dict[VtfPF, tuple[int,] | str] = {
        VtfPF.DXT1: (1,),
        VtfPF.DXT3: (3,),
        VtfPF.DXT5: (5,),
        VtfPF.RGB888: "RGB",
        VtfPF.BGR888: "BGR",
        VtfPF.RGBA8888: "RGBA",
        VtfPF.A8: "A",
        VtfPF.I8: "L",
        VtfPF.IA88: "LA",
        VtfPF.UV88: "RG",
    }
    try:
        encoder_args = args[pixel_format]
    except KeyError:
        msg = f"Unsupported pixel format: {pixel_format!r}"
        raise VTFException(msg)

    codec_name = "bcn" if pixel_format in BLOCK_COMPRESSED else "raw"
    tile = [ImageFile._Tile(codec_name, (0, 0) + im.size, fp.tell(), encoder_args)]
    ImageFile._save(im, fp, tile, _get_texture_size(pixel_format, *im.size))


def _closest_power(x: int) -> int:
    possible_results = round(log(x, 2)), ceil(log(x, 2))
    return 2 ** min(possible_results, key=lambda z: abs(x - 2**z))


class VtfImageFile(ImageFile.ImageFile):
    format = "VTF"
    format_description = "Valve Texture Format"

    def _open(self):
        if not _accept(self.fp.read(12)):
            msg = "not a VTF file"
            raise SyntaxError(msg)

        self.fp.seek(4)
        version = struct.unpack("<2I", self.fp.read(8))
        if version > (7, 4):
            msg = f"Unsupported VTF version: {version}"
            raise VTFException(msg)

        header = VTFHeader(
            *struct.unpack(HEADER_V70, self.fp.read(struct.calcsize(HEADER_V70)))
        )
        self.fp.seek(header.header_size)

        pixel_format = header.pixel_format
        if pixel_format in (
            VtfPF.DXT1_ONEBITALPHA,
            VtfPF.DXT1,
            VtfPF.DXT3,
            VtfPF.DXT5,
            VtfPF.RGBA8888,
            VtfPF.BGRA8888,
            VtfPF.A8,
        ):
            self._mode = "RGBA"
        elif pixel_format in (VtfPF.RGB888, VtfPF.BGR888, VtfPF.UV88):
            self._mode = "RGB"
        elif pixel_format == VtfPF.I8:
            self._mode = "L"
        elif pixel_format == VtfPF.IA88:
            self._mode = "LA"
        else:
            msg = f"Unsupported VTF pixel format: {pixel_format}"
            raise VTFException(msg)

        if pixel_format in BLOCK_COMPRESSED:
            codec_name = "bcn"
            if pixel_format in (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA):
                args = (1, "DXT1")
            elif pixel_format == VtfPF.DXT3:
                args = (2, "DXT3")
            else:
                args = (3, "DXT5")
        else:
            codec_name = "raw"
            try:
                args = {
                    VtfPF.RGBA8888: "RGBA",
                    VtfPF.RGB888: "RGB",
                    VtfPF.BGR888: "BGR",
                    VtfPF.BGRA8888: "BGRA",
                    VtfPF.UV88: "RG",
                    VtfPF.I8: "L",
                    VtfPF.A8: "A",
                    VtfPF.IA88: "LA",
                }[pixel_format]
            except KeyError:
                msg = f"Unsupported VTF pixel format: {pixel_format}"
                raise VTFException(msg)

        self._size = (header.width, header.height)

        data_start = self.fp.tell()
        data_start += _get_texture_size(
            header.low_pixel_format, header.low_width, header.low_height
        )
        min_res = 4 if pixel_format in BLOCK_COMPRESSED else 1
        for mip_id in range(header.mipmap_count - 1, 0, -1):
            mip_width = max(header.width >> mip_id, min_res)
            mip_height = max(header.height >> mip_id, min_res)

            data_start += _get_texture_size(pixel_format, mip_width, mip_height)

        self.tile = [ImageFile._Tile(codec_name, (0, 0) + self.size, data_start, args)]


def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
    if im.mode not in ("RGB", "RGBA", "L", "LA"):
        msg = f"cannot write mode {im.mode} as VTF"
        raise OSError(msg)

    encoderinfo = im.encoderinfo
    version = encoderinfo.get("version", (7, 4))
    if version > (7, 4):
        msg = f"Unsupported VTF version: {version}"
        raise VTFException(msg)

    pixel_format = (
        encoderinfo.get("pixel_format")
        or {
            "RGB": VtfPF.RGB888,
            "RGBA": VtfPF.RGBA8888,
            "L": VtfPF.I8,
            "LA": VtfPF.IA88,
        }[im.mode]
    )
    generate_mips = encoderinfo.get("generate_mips", True)

    flags = CompiledVtfFlags(0)

    if pixel_format in (
        VtfPF.DXT3,
        VtfPF.DXT5,
        VtfPF.RGBA8888,
        VtfPF.BGRA8888,
        VtfPF.A8,
        VtfPF.IA88,
    ):
        flags |= CompiledVtfFlags.EIGHTBITALPHA
    im = im.resize((_closest_power(im.width), _closest_power(im.height)))
    width, height = im.size

    mipmap_count = _get_mipmap_count(width, height) if generate_mips else 0

    thumb = im.convert("RGB")
    thumb.thumbnail((min(16, width), min(16, height)))
    thumb = thumb.resize((_closest_power(thumb.width), _closest_power(thumb.height)))

    header = VTFHeader(
        0,
        width,
        height,
        flags,
        1,
        0,
        1.0,
        1.0,
        1.0,
        1.0,
        pixel_format,
        mipmap_count,
        VtfPF.DXT1,
        thumb.width,
        thumb.height,
        1,
        2,
    )

    fp.write(b"VTF\x00" + struct.pack("<2I", *version))
    if version < (7, 2):
        size = struct.calcsize(HEADER_V70) + 12
        header = header._replace(header_size=size + (16 - size % 16))
        fp.write(struct.pack(HEADER_V70, *header[:15]))
    elif version == (7, 2):
        size = struct.calcsize(HEADER_V72) + 12
        header = header._replace(header_size=size + (16 - size % 16))
        fp.write(struct.pack(HEADER_V72, *header[:16]))
    else:
        size = struct.calcsize(HEADER_V73) + 12
        header = header._replace(header_size=size + (16 - size % 16))
        fp.write(struct.pack(HEADER_V73, *header))

    if version > (7, 2):
        fp.write(b"\x01\x00\x00\x00")
        fp.write(struct.pack("<I", header.header_size))
        fp.write(b"\x30\x00\x00\x00")
        fp.write(
            struct.pack(
                "<I",
                header.header_size
                + _get_texture_size(VtfPF.DXT1, thumb.width, thumb.height),
            )
        )
    else:
        fp.write(b"\x00" * (16 - fp.tell() % 16))
    _write_image(fp, thumb, VtfPF.DXT1)

    min_size = 4 if pixel_format in BLOCK_COMPRESSED else 1

    for mip_id in range(mipmap_count - 1, 0, -1):
        mip_width = max(min_size, width >> mip_id)
        mip_height = max(min_size, height >> mip_id)
        mip = im.resize((mip_width, mip_height))
        _write_image(fp, mip, pixel_format)
    _write_image(fp, im, pixel_format)


def _accept(prefix: bytes) -> bool:
    valid_header = prefix.startswith(b"VTF\x00")
    valid_version = struct.unpack_from("<2I", prefix, 4) >= (7, 0)
    return valid_header and valid_version


Image.register_open(VtfImageFile.format, VtfImageFile, _accept)
Image.register_save(VtfImageFile.format, _save)
Image.register_extension(VtfImageFile.format, ".vtf")