import json
import struct
import typing  # noqa

from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from shlex import shlex
from typing import BinaryIO, Union

# https://learn.microsoft.com/en-us/windows/win32/debug/pe-format
endian = "<"
IMAGE_DIRECTORY_ENTRY = 2
VERSION_RESOURCE = 16
RESOURCE_NODE = ".rsrc"
PE32_x86 = "0x10b"
PE32_x64 = "0x20b"
VS_VERSION_INFO_MAGIC = "0xfeef04bd0000"
VS_VERSION_INFO_ID = "VS_VERSION_INFO"


class VersionMatch(Enum):
    LOCAL_OLDER = 1
    LOCAL_NEWER = 2
    SAME_VERSION = 3


class u8:
    fmt = "B"


class u16:
    fmt = "H"


class u32:
    fmt = "L"


class u64:
    fmt = "Q"


class i8:
    fmt = "b"


class i16:
    fmt = "h"


class i32:
    fmt = "l"


class i64:
    fmt = "q"


class PackedData:
    @classmethod
    def unpack(cls, data: BinaryIO):
        r = []
        for key, value in cls.__annotations__.items():
            if value == str:
                f = data.read(8).rstrip(b"\x00\x00").decode()
            else:
                fmt = endian + (value.fmt)
                size = struct.calcsize(fmt)
                f = struct.unpack(fmt, data.read(size))[0]
            r.append(f)
        return cls(*r)


@dataclass(slots=True, frozen=True)
class COFF_FILE_HDR(PackedData):
    machine_type: u16
    number_of_sections: u16
    timestamp: u32
    pointer_to_symbol_table: u32
    number_of_symbols: u32
    size_of_optional_header: u16
    characteristics: u16


@dataclass(slots=True, frozen=True)
class OPTIONAL_HDR_X86(PackedData):
    magic: u16
    major_linker_ver: u8
    minor_linker_ver: u8
    size_of_code: u32
    size_of_initialized_data: u32
    size_of_uninitialized_data: u32
    address_of_entry_point: u32
    base_of_code: u32
    base_of_data: u32


@dataclass(slots=True, frozen=True)
class OPTIONAL_HDR_X64(PackedData):
    magic: u16
    major_linker_ver: u8
    minor_linker_ver: u8
    size_of_code: u32
    size_of_initialized_data: u32
    size_of_uninitialized_data: u32
    address_of_entry_point: u32
    base_of_code: u32


@dataclass(slots=True, frozen=True)
class OPTIONAL_HDR_WIN_X86(PackedData):
    image_base: u32
    section_alignment: u32
    file_alignment: u32
    major_operating_system_version: u16
    minor_operating_system_version: u16
    major_image_version: u16
    minor_image_version: u16
    major_subsystem_version: u16
    minor_subsystem_version: u16
    win32_version_value: u32
    size_of_image: u32
    size_of_headers: u32
    checksum: u32
    subsystem: u16
    dll_characteristics: u16
    size_of_stack_reserve: u32
    size_of_stack_commit: u32
    size_of_heap_reserve: u32
    size_of_heap_commit: u32
    loader_flags: u32
    number_of_rva_and_sizes: u32


@dataclass(slots=True, frozen=True)
class OPTIONAL_HDR_WIN_X64(PackedData):
    image_base: u64
    section_alignment: u32
    file_alignment: u32
    major_operating_system_version: u16
    minor_operating_system_version: u16
    major_image_version: u16
    minor_image_version: u16
    major_subsystem_version: u16
    minor_subsystem_version: u16
    win32_version_value: u32
    size_of_image: u32
    size_of_headers: u32
    checksum: u32
    subsystem: u16
    dll_characteristics: u16
    size_of_stack_reserve: u64
    size_of_stack_commit: u64
    size_of_heap_reserve: u64
    size_of_heap_commit: u64
    loader_flags: u32
    number_of_rva_and_sizes: u32


@dataclass(slots=True, frozen=True)
class DATA_DIR(PackedData):
    virtual_address: u32
    size: u32


@dataclass(slots=True, frozen=True)
class SECTION_HDR(PackedData):
    name: str
    virtual_size: u32
    virtual_address: u32
    size_of_raw_data: u32
    pointer_to_raw_data: u32
    pointer_to_relocations: u32
    pointer_to_line_numbers: u32
    number_of_relocations: u16
    number_of_line_numbers: u16
    characteristics: u32


@dataclass(slots=True, frozen=True)
class RESOURCE_DIRECTORY_TABLE(PackedData):
    characteristics: u32
    timestamp: u32
    major_version: u16
    minor_version: u16
    number_of_name_entries: u16
    number_of_id_entries: u16


@dataclass(slots=True, frozen=True)
class RESOURCE_DIRECTORY_ENTRY(PackedData):
    """
    This field is either a string identifying a data leaf
    (if the high bit is set) or an ID to another nested directory
    (if the high bit is clear). The outermost level is always a
    directory. If it is a name, the lower 31 bits are the offset from the
    beginning of the resource section's raw data to the name
    (the name consists of 16 bits length and trailing wide characters,
    in Unicode, not 0-terminated).
    """

    name_or_id: u32
    data_or_subdir: u32


@dataclass(slots=True, frozen=True)
class RESOURCE_DATA_ENTRY(PackedData):
    data_rva: u32
    size: u32
    codepage: u32
    reserved: u32


@dataclass(slots=True, frozen=True)
class VS_VERSION_INFO_HDR(PackedData):
    size: u16
    value_length: u16
    value_type: u16


@dataclass(slots=True, frozen=True)
class DayZVersion:
    major: int
    minor: int
    patch: int


@dataclass(slots=True, frozen=True)
class FileVersion:
    major: int
    minor: int
    build: int
    revision: int


@dataclass(slots=True, frozen=True)
class Result:
    local: Union[str, None]
    remote: str
    build: str
    path: Union[Path, None]
    match: VersionMatch
    error: Union[Exception, None]


class PeFileError(Exception):
    """Expected contents missing from headers or resource nodes"""

    pass


class AppNotInstalledError(Exception):
    """App not present in user's libraryfolders"""

    pass


class AppMovedError(Exception):
    """VDF points to a nonexistent location on disk"""

    pass


class VDFLoadError(Exception):
    """Malformed VDF or JSON conversion"""

    pass


def parse_version_number(data: BinaryIO):
    # https://learn.microsoft.com/en-us/windows/win32/api/verrsrc/ns-verrsrc-vs_fixedfileinfo
    minor = struct.unpack("<L", data.read(4))[0] >> 16 & 0xFFFF
    major = struct.unpack("<L", data.read(4))[0] >> 0 & 0xFFFF
    build = struct.unpack("<L", data.read(4))[0] >> 0 & 0xFFFF
    revision = struct.unpack("<L", data.read(4))[0] >> 16 & 0xFFFF
    return FileVersion(major, minor, build, revision)


def seek_to_hex(address: str, data: BinaryIO) -> None:
    data.seek(int(address, base=16))


def seek_to_pe_stub(data: BinaryIO) -> None:
    MAGIC = "0x3c"
    seek_to_hex(MAGIC, data)
    e_lfanew = hex(struct.unpack("<h", data.read(2))[0])
    seek_to_hex(e_lfanew, data)
    pe_stub = data.read(4).rstrip(b"\x00\x00\x00\x00").decode()
    if pe_stub != "PE":
        raise PeFileError("missing PE header data")


def get_dayz_version(file: Path) -> DayZVersion | Exception:
    try:
        version = get_version(file)
    except Exception as e:
        raise e
    patch = str(version.build) + str(version.revision)
    dz_vers = DayZVersion(version.major, version.minor, int(patch))
    return dz_vers


def dayz_version_to_str(v: DayZVersion) -> str:
    return ".".join(str(el) for el in [v.major, v.minor, v.patch])


def dayz_version_from_str(v: str) -> DayZVersion:
    vers = v.split(".")
    assert len(vers) == 3
    return DayZVersion(*[int(el) for el in vers])


def get_version(file):
    with open(file, "rb") as f:
        seek_to_pe_stub(f)
        COFF = COFF_FILE_HDR.unpack(f)

        pos = f.tell()
        blob = f.read()
        magic = hex(struct.unpack("<H", (blob[0:2]))[0])
        f.seek(pos)

        if magic == PE32_x86:
            OPTIONAL_HDR_X86.unpack(f)
            OBJW = OPTIONAL_HDR_WIN_X86.unpack(f)
        else:
            OPTIONAL_HDR_X64.unpack(f)
            OBJW = OPTIONAL_HDR_WIN_X64.unpack(f)

        if OBJW.number_of_rva_and_sizes < 1:
            raise PeFileError("no data resource directory")

        data_dirs = []
        for rva in range(OBJW.number_of_rva_and_sizes):
            data_dir = DATA_DIR.unpack(f)
            data_dirs.append(data_dir)

        res_dir = data_dirs[IMAGE_DIRECTORY_ENTRY]
        dir_va = res_dir.virtual_address

        for section in range(COFF.number_of_sections):
            hdr = SECTION_HDR.unpack(f)
            if hdr.name == RESOURCE_NODE:
                va = hdr.virtual_address
                ptr = hdr.pointer_to_raw_data
                offset = dir_va - va + ptr
                seek_to_hex(hex(offset), f)
                break

        if hdr.name != RESOURCE_NODE:
            raise PeFileError("no root resource node found")

        table = RESOURCE_DIRECTORY_TABLE.unpack(f)
        total = table.number_of_name_entries + table.number_of_id_entries

        for entry in range(total):
            entry = RESOURCE_DIRECTORY_ENTRY.unpack(f)
            if entry.name_or_id == VERSION_RESOURCE:
                while entry.data_or_subdir & (1 << 31):
                    shift = entry.data_or_subdir & ~(1 << 31)
                    seek_to_hex(hex(offset + shift), f)
                    table = RESOURCE_DIRECTORY_TABLE.unpack(f)
                    total = (
                        table.number_of_name_entries
                        + table.number_of_id_entries
                    )
                    for entry in range(total):
                        entry = RESOURCE_DIRECTORY_ENTRY.unpack(f)
                break
            if entry.name_or_id > VERSION_RESOURCE:
                raise PeFileError("no version info node found")
        seek_to_hex(hex(offset + entry.data_or_subdir), f)

        data = RESOURCE_DATA_ENTRY.unpack(f)
        # https://stackoverflow.com/questions/2170843/va-virtual-address-rva-relative-virtual-address
        offset = data.data_rva - hdr.virtual_address + hdr.pointer_to_raw_data
        seek_to_hex(hex(offset), f)

        hdr = VS_VERSION_INFO_HDR.unpack(f)
        # https://learn.microsoft.com/en-us/windows/win32/menurc/vs-versioninfo
        byte_len = len(VS_VERSION_INFO_ID.encode("utf-16le"))
        label = f.read(byte_len).decode("utf-16le")
        if label != VS_VERSION_INFO_ID:
            raise PeFileError(f"header identifier != '{VS_VERSION_INFO_ID}'")
        f.read(32 - byte_len)

        identifier = hex(struct.unpack("<Q", f.read(8))[0])
        if identifier != VS_VERSION_INFO_MAGIC:
            raise PeFileError(
                f"{VS_VERSION_INFO_ID} address != '{VS_VERSION_INFO_MAGIC}'"
            )

        try:
            version = parse_version_number(f)
        except Exception as e:
            return PeFileError(e)

        return version


def get_pefile_path(path: str, appid: int) -> Path:
    binary = "DayZ_x64.exe"
    identifier = {221100: "DayZ", 1024020: "DayZ Exp"}
    name = identifier[appid]

    pe_path = None
    path = path + "/steamapps/libraryfolders.vdf"

    with open(path, "r") as f:
        try:
            j = json.loads(vdf_to_json(f))
        except Exception:
            raise VDFLoadError("Failed to parse libraryfolders")

        for obj in j["libraryfolders"]:
            if str(appid) in j["libraryfolders"][obj]["apps"]:
                pe_path = j["libraryfolders"][obj]["path"]
                pe_path += f"/steamapps/common/{name}/{binary}"
                break

    if pe_path is None:
        raise AppNotInstalledError(
            f"Failed to find a libraryfolder for the appid '{appid}'"
        )

    pe_path = Path(pe_path)
    if pe_path.exists() is False:
        raise AppMovedError(
            f"Path '{pe_path}' specified in libraryfolders does not exist"
        )

    return pe_path


def compare_versions(local: DayZVersion, remote: DayZVersion):
    """
    packaging.version module is not available OOTB on some distributions
    """
    if dayz_version_to_str(local) == dayz_version_to_str(remote):
        return VersionMatch.SAME_VERSION

    if local.major < remote.major:
        return VersionMatch.LOCAL_OLDER
    if local.major > remote.major:
        return VersionMatch.LOCAL_NEWER
    if local.major == remote.major:
        if local.minor < remote.minor:
            return VersionMatch.LOCAL_OLDER
        if local.minor > remote.minor:
            return VersionMatch.LOCAL_NEWER
        if local.minor == remote.minor:
            if local.patch < remote.patch:
                return VersionMatch.LOCAL_OLDER
            if local.patch > remote.patch:
                return VersionMatch.LOCAL_NEWER
            if local.patch == remote.patch:
                return VersionMatch.SAME_VERSION


def vdf_to_json(stream):
    def _istr(indent, string):
        return (indent * "  ") + string

    jbuf = "{\n"
    lex = shlex(stream)
    indent = 1

    while True:
        tok = lex.get_token()
        if not tok:
            return jbuf + "}\n"
        if tok == "}":
            indent -= 1
            jbuf += _istr(indent, "}")
            ntok = lex.get_token()
            lex.push_token(ntok)
            if ntok and ntok != "}":
                jbuf += ","
            jbuf += "\n"
        else:
            ntok = lex.get_token()
            if ntok == "{":
                jbuf += _istr(indent, tok + ": {\n")
                indent += 1
            else:
                jbuf += _istr(indent, tok + ": " + ntok)
                ntok = lex.get_token()
                lex.push_token(ntok)
                if ntok != "}":
                    jbuf += ","
                jbuf += "\n"