from __future__ import annotations

import builtins
import re
import sys
from io import BufferedReader, IOBase
from struct import Struct, unpack
from types import MethodType
from typing import Any, Callable, ClassVar, Dict, List, Literal, Optional, Tuple, Union

reNot0 = re.compile(b"(.*?)\x00", re.S)

SYS_ENDIAN = "<" if sys.byteorder == "little" else ">"
Endianess = Literal["<", ">"]

# generate unpack and unpack_from functions
TYPE_PARAM_SIZE_LIST = [
    ("short", "h"),
    ("u_short", "H"),
    ("int", "i"),
    ("u_int", "I"),
    ("long", "q"),
    ("u_long", "Q"),
    ("half", "e"),
    ("float", "f"),
    ("double", "d"),
]

MEMORY_FUNCTIONS: Dict[Endianess, Dict[str, Callable[["EndianBinaryReader_Memoryview"], Any]]] = {"<": {}, ">": {}}
STREAM_FUNCTIONS: Dict[Endianess, Dict[str, Callable[["EndianBinaryReader_Streamable"], Any]]] = {"<": {}, ">": {}}


class EndianBinaryReader:
    Length: int
    Position: int
    BaseOffset: int
    _endian: Endianess
    _function_map: ClassVar[Dict[Endianess, Dict[str, Callable]]]

    def __new__(
        cls,
        item: Union[bytes, bytearray, memoryview, IOBase, str],
        endian: Endianess = ">",
        offset: int = 0,
    ):
        if isinstance(item, (bytes, bytearray, memoryview)):
            obj = super(EndianBinaryReader, cls).__new__(EndianBinaryReader_Memoryview)  # type: ignore
        elif isinstance(item, IOBase):
            obj = super(EndianBinaryReader, cls).__new__(EndianBinaryReader_Streamable)  # type: ignore
        elif isinstance(item, str):
            obj = super(EndianBinaryReader, cls).__new__(EndianBinaryReader_Streamable_LocalFile)  # type: ignore
        elif isinstance(item, EndianBinaryReader):
            item = item.stream if isinstance(item, EndianBinaryReader_Streamable) else item.view
            return EndianBinaryReader(item, endian, offset)
        elif hasattr(item, "read"):
            if hasattr(item, "seek") and hasattr(item, "tell"):
                obj = super(EndianBinaryReader, cls).__new__(EndianBinaryReader_Streamable)
            else:
                item = item.read()
                obj = super(EndianBinaryReader, cls).__new__(EndianBinaryReader_Memoryview)
        else:
            raise TypeError("Unsupported type for EndianBinaryReader: %s" % type(item))
        return obj

    def __init__(self, item, endian: Endianess = ">", offset: int = 0):
        self._endian = ""  # type: ignore
        self.endian = endian
        self.BaseOffset = offset
        self.Position = 0

    @property
    def endian(self) -> Endianess:
        return self._endian

    @endian.setter
    def endian(self, value: Endianess):
        if value not in ("<", ">"):
            raise ValueError("Invalid endian")
        if value != self._endian:
            for func_name, func in self._function_map[value].items():
                setattr(self, func_name, MethodType(func, self))
            self._endian = value

    @property
    def bytes(self) -> builtins.bytes:
        # implemented by Streamable and Memoryview versions
        return b""

    def read(self, size: Optional[int] = -1, /) -> builtins.bytes:
        # implemented by Streamable and Memoryview versions
        return b""

    def read_byte(self) -> int:
        return unpack(self.endian + "b", self.read(1))[0]

    def read_u_byte(self) -> int:
        return unpack(self.endian + "B", self.read(1))[0]

    def read_bytes(self, num: int) -> builtins.bytes:
        return self.read(num)

    def read_short(self) -> int:
        return unpack(self.endian + "h", self.read(2))[0]

    def read_int(self) -> int:
        return unpack(self.endian + "i", self.read(4))[0]

    def read_long(self) -> int:
        return unpack(self.endian + "q", self.read(8))[0]

    def read_u_short(self) -> int:
        return unpack(self.endian + "H", self.read(2))[0]

    def read_u_int(self) -> int:
        return unpack(self.endian + "I", self.read(4))[0]

    def read_u_long(self) -> int:
        return unpack(self.endian + "Q", self.read(8))[0]

    def read_float(self) -> float:
        return unpack(self.endian + "f", self.read(4))[0]

    def read_double(self) -> float:
        return unpack(self.endian + "d", self.read(8))[0]

    def read_boolean(self) -> bool:
        return bool(unpack(self.endian + "?", self.read(1))[0])

    def read_string(self, size: Optional[int] = None) -> str:
        if size is None:
            return self.read_string_to_null()
        else:
            raw = self.read_bytes(size)
            return raw.decode("utf8", "surrogateescape")

    def read_string_to_null(self, max_length: int = 32767) -> str:
        ret = []
        c = b""
        while c != b"\0" and len(ret) < max_length and self.Position != self.Length:
            ret.append(c)
            c = self.read(1)
            if not c:
                raise ValueError("Unterminated string: %r" % ret)
        return b"".join(ret).decode("utf8", "surrogateescape")

    def read_aligned_string(self) -> str:
        length = self.read_int()
        if 0 < length <= self.Length - self.Position:
            string_data = bytes(self.read_bytes(length))
            result = string_data.decode("utf8", "surrogateescape")
            self.align_stream()
            return result
        return ""

    def align_stream(self, alignment=4):
        self.Position += (alignment - self.Position % alignment) % alignment

    def read_byte_array(self) -> builtins.bytes:
        return self.read(self.read_int())

    def read_array(self, command: Callable, length: int) -> list:
        return [command() for _ in range(length)]

    def read_array_struct(self, param: str, length: Optional[int] = None) -> tuple:
        if length is None:
            length = self.read_int()
        struct = Struct(f"{self.endian}{length}{param}")
        return struct.unpack(self.read(struct.size))

    def read_boolean_array(self, length: Optional[int] = None) -> Tuple[bool, ...]:
        return self.read_array_struct("?", length)

    def read_u_byte_array(self, length: Optional[int] = None) -> Tuple[int, ...]:
        return self.read_array_struct("B", length)

    def read_u_short_array(self, length: Optional[int] = None) -> Tuple[int, ...]:
        return self.read_array_struct("h", length)

    def read_short_array(self, length: Optional[int] = None) -> Tuple[int, ...]:
        return self.read_array_struct("H", length)

    def read_int_array(self, length: Optional[int] = None) -> Tuple[int, ...]:
        return self.read_array_struct("i", length)

    def read_u_int_array(self, length: Optional[int] = None) -> Tuple[int, ...]:
        return self.read_array_struct("I", length)

    def read_long_array(self, length: Optional[int] = None) -> Tuple[int, ...]:
        return self.read_array_struct("q", length)

    def read_u_long_array(self, length: Optional[int] = None) -> Tuple[int, ...]:
        return self.read_array_struct("Q", length)

    def read_float_array(self, length: Optional[int] = None) -> Tuple[float, ...]:
        return self.read_array_struct("f", length)

    def read_double_array(self, length: Optional[int] = None) -> Tuple[float, ...]:
        return self.read_array_struct("d", length)

    def read_string_array(self) -> List[str]:
        return self.read_array(self.read_aligned_string, self.read_int())

    def real_offset(self) -> int:
        """Returns offset in the underlying file.
        (Not working with unpacked streams.)
        """
        return self.BaseOffset + self.Position

    def read_the_rest(self, obj_start: int, obj_size: int) -> builtins.bytes:
        """Returns the rest of the current reader bytes."""
        return self.read_bytes(obj_size - (self.Position - obj_start))

    def seek(self, offset: int, whence: int = 0) -> int:
        if whence == 0:
            new_pos = offset
        elif whence == 1:
            new_pos = self.Position + offset
        elif whence == 2:
            new_pos = self.Length + offset
        else:
            raise ValueError("Invalid whence value")
        if new_pos < 0:
            raise ValueError("New position is before the start of the stream")
        self.Position = new_pos
        return self.Position

    def tell(self) -> int:
        return self.Position


class EndianBinaryReader_Memoryview(EndianBinaryReader):
    __slots__ = ("view", "_endian", "BaseOffset", "Position", "Length")
    _endian: Endianess
    view: memoryview
    _function_map = MEMORY_FUNCTIONS

    def __init__(self, view, endian: Endianess = ">", offset: int = 0):
        super().__init__(view, endian=endian, offset=offset)
        self.view = memoryview(view)
        self.Length = len(view)

    @property
    def bytes(self):
        return self.view.tobytes()

    def dispose(self) -> None:
        self.view.release()

    def read(self, size: Optional[int] = -1, /):
        if not size:
            return b""
        if size == -1:
            size = self.Length - self.Position
        ret = self.view[self.Position : self.Position + size]
        self.Position += size
        return ret.tobytes()

    def read_array_struct(self, param: str, length: Optional[int] = None) -> tuple:
        if length is None:
            length = self.read_int()
        struct = Struct(f"{self.endian}{length}{param}")
        value = struct.unpack_from(self.view, self.Position)
        self.Position += struct.size
        return value

    def read_aligned_string(self) -> str:
        length = self.read_int()
        if 0 < length <= self.Length - self.Position:
            string_data = self.read_bytes(length)
            result = bytes(string_data).decode("utf8", "surrogateescape")
            self.align_stream()
            return result
        return ""

    def read_string_to_null(self, max_length: int = 32767) -> str:
        match = reNot0.search(self.view, self.Position, self.Position + max_length)
        if not match:
            if self.Position + max_length >= self.Length:
                raise Exception("String not terminated")
            else:
                return bytes(self.read_bytes(max_length)).decode("utf8", "surrogateescape")
        ret = match[1].decode("utf8", "surrogateescape")
        self.Position = match.end()
        return ret


class EndianBinaryReader_Streamable(EndianBinaryReader):
    __slots__ = ("stream", "_endian", "BaseOffset")
    stream: BufferedReader
    _function_map = STREAM_FUNCTIONS

    def __init__(self, stream: BufferedReader, endian: Endianess = ">", offset: int = 0):
        self.stream = stream
        super().__init__(stream, endian=endian, offset=offset)
        self.read = self.stream.read

    @property
    def Position(self) -> int:
        return self.stream.tell() - self.BaseOffset

    @Position.setter
    def Position(self, value: int):
        if value < 0:
            raise ValueError("Position cannot be negative")
        self.stream.seek(value + self.BaseOffset)

    @property
    def Length(self):  # type: ignore
        pos = self.Position
        length = self.stream.seek(0, 2) - self.BaseOffset
        self.Position = pos
        return length

    @property
    def bytes(self):
        last_pos = self.Position
        self.Position = 0
        ret = self.read(self.Length)
        self.Position = last_pos  # type: ignore
        return ret

    def dispose(self):
        self.stream.close()
        pass


class EndianBinaryReader_Streamable_LocalFile(EndianBinaryReader_Streamable):
    def __init__(self, path: str, endian: Endianess = ">", offset: int = 0):
        super().__init__(open(path, "rb"), endian=endian, offset=offset)

    def __del__(self):
        self.stream.close()


for endian_s in ("<", ">"):
    for reader_type_name, struct_type_char in TYPE_PARAM_SIZE_LIST:
        func_name = f"read_{reader_type_name}"
        struct = Struct(f"{endian_s}{struct_type_char}")

        def memory_read_func(self: EndianBinaryReader_Memoryview, /, struct=struct):
            value = struct.unpack_from(self.view, self.Position)[0]
            self.Position += struct.size
            return value

        def stream_read_func(self: EndianBinaryReader_Streamable, /, struct=struct):
            return struct.unpack(self.stream.read(struct.size))[0]

        MEMORY_FUNCTIONS[endian_s][func_name] = memory_read_func
        STREAM_FUNCTIONS[endian_s][func_name] = stream_read_func