from __future__ import annotations

import binascii
import re
import struct
from typing import Any

from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql import functions
from sqlalchemy.sql.functions import FunctionElement
from sqlalchemy.types import to_instance

from geoalchemy2.exc import ArgumentError

BinasciiError = binascii.Error

function_registry: set[str] = set()


class _SpatialElement:
    """The base class for public spatial elements.

    Args:
        data: The first argument passed to the constructor is the data wrapped
            by the ``_SpatialElement`` object being constructed.
        srid: An integer representing the spatial reference system. E.g. ``4326``.
            Default value is ``-1``, which means no/unknown reference system.
        extended: A boolean indicating whether the extended format (EWKT or EWKB)
            is used. Default is ``None``.

    """

    # Define __slots__ to restrict attributes in this class.
    # This is done intentionally to improve performance by preventing
    # the creation of a dynamic __dict__ for each instance.
    __slots__ = ("srid", "data", "extended")

    def __init__(self, data, srid: int = -1, extended: bool | None = None) -> None:
        self.srid = srid
        self.data = data
        self.extended = extended

    def __str__(self) -> str:
        return self.desc

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__} at 0x{id(self):x}; {self}>"  # pragma: no cover

    def __eq__(self, other) -> bool:
        try:
            return (
                self.extended == other.extended
                and self.srid == other.srid
                and self.desc == other.desc
            )
        except AttributeError:
            return False

    def __ne__(self, other) -> bool:
        return not self.__eq__(other)

    def __hash__(self):
        return hash((self.desc, self.srid, self.extended))

    def __getattr__(self, name):
        #
        # This is how things like lake.geom.ST_Buffer(2) creates
        # SQL expressions of this form:
        #
        # ST_Buffer(ST_GeomFromWKB(:ST_GeomFromWKB_1), :param_1)
        #

        # Raise an AttributeError when the attribute name doesn't start
        # with st_. This is to be nice with other libraries that use
        # some ducktyping (e.g. hasattr(element, "copy")) to determine
        # the type of the element.

        if name.lower() not in function_registry:
            raise AttributeError

        # We create our own _FunctionGenerator here, and use it in place of
        # SQLAlchemy's "func" object. This is to be able to "bind" the
        # function to the SQL expression. See also GenericFunction above.
        func_ = functions._FunctionGenerator(expr=self)
        return getattr(func_, name)

    def __getstate__(self) -> dict[str, Any]:
        state = {
            "srid": self.srid,
            "data": str(self),
            "extended": self.extended,
        }
        return state

    def __setstate__(self, state: dict[str, Any]) -> None:
        self.srid = state["srid"]
        self.extended = state["extended"]
        self.data = self._data_from_desc(state["data"])

    @staticmethod
    def _data_from_desc(desc):
        raise NotImplementedError()  # pragma: no cover


class WKTElement(_SpatialElement):
    """Instances of this class wrap a WKT or EWKT value.

    Usage examples::

        wkt_element_1 = WKTElement('POINT(5 45)')
        wkt_element_2 = WKTElement('POINT(5 45)', srid=4326)
        wkt_element_3 = WKTElement('SRID=4326;POINT(5 45)', extended=True)

    Note::
        This class uses ``__slots__`` to restrict its attributes and improve memory efficiency by
        preventing the creation of a dynamic ``__dict__`` for each instance.
        If you require dynamic attributes or support for weak references, use the
        ``DynamicWKTElement`` subclass, which provides these capabilities.
    """

    __slots__ = ()

    _REMOVE_SRID = re.compile("(SRID=([0-9]+); ?)?(.*)")
    SPLIT_WKT_PATTERN = re.compile(r"((SRID=\d+) *; *)?([\w ]+) *(\([-\d\. ,\(\)eE]+\))")

    geom_from: str = "ST_GeomFromText"
    geom_from_extended_version: str = "ST_GeomFromEWKT"

    def __init__(self, data: str, srid: int = -1, extended: bool | None = None) -> None:
        if extended is None:
            extended = data.startswith("SRID=")
        if extended and srid == -1:
            # read srid from EWKT
            data_s = data.split(";")
            if len(data_s) != 2:
                raise ArgumentError(f"invalid EWKT string {data}")
            header = data_s[0]
            try:
                srid = int(header[5:])
            except ValueError:
                raise ArgumentError(f"invalid EWKT string {data}") from None
        _SpatialElement.__init__(self, data, srid, extended)

    @property
    def desc(self) -> str:
        """This element's description string."""
        return self.data

    @staticmethod
    def _data_from_desc(desc):
        return desc

    def as_wkt(self) -> WKTElement:
        if self.extended:
            srid_match = self._REMOVE_SRID.match(self.data)
            assert srid_match is not None
            return WKTElement(srid_match.group(3), self.srid, extended=False)
        return WKTElement(self.data, self.srid, self.extended)

    def as_ewkt(self) -> WKTElement:
        if not self.extended and self.srid != -1:
            data = f"SRID={self.srid};" + self.data
            return WKTElement(data, extended=True)
        return WKTElement(self.data, self.srid, self.extended)


class DynamicWKTElement(WKTElement):
    """This is a subclass of ``WKTElement`` that allows dynamic attributes.

    It is useful when you need to add attributes dynamically to the object.
    """

    __slots__ = ("__dict__", "__weakref__")


class WKBElement(_SpatialElement):
    """Instances of this class wrap a WKB or EWKB value.

    Geometry values read from the database are converted to instances of this
    type. In most cases you won't need to create ``WKBElement`` instances
    yourself.

    If ``extended`` is ``True`` and ``srid`` is ``-1`` at construction time
    then the SRID will be read from the EWKB data.

    Note: you can create ``WKBElement`` objects from Shapely geometries
    using the :func:`geoalchemy2.shape.from_shape` function.

    Note::
        This class uses ``__slots__`` to restrict its attributes and improve memory efficiency by
        preventing the creation of a dynamic ``__dict__`` for each instance.
        If you require dynamic attributes or support for weak references, use the
        ``DynamicWKBElement`` subclass, which provides these capabilities.
    """

    __slots__ = ()

    geom_from: str = "ST_GeomFromWKB"
    geom_from_extended_version: str = "ST_GeomFromEWKB"

    def __init__(
        self, data: str | bytes | memoryview, srid: int = -1, extended: bool | None = None
    ) -> None:
        if srid == -1 or extended is None or extended:
            # read srid from the EWKB
            #
            # WKB struct {
            #    byte    byteOrder;
            #    uint32  wkbType;
            #    uint32  SRID;
            #    struct  geometry;
            # }
            # byteOrder enum {
            #     WKB_XDR = 0,  // Most Significant Byte First
            #     WKB_NDR = 1,  // Least Significant Byte First
            # }
            # See https://trac.osgeo.org/postgis/browser/branches/3.0/doc/ZMSgeoms.txt
            # for more details about WKB/EWKB specifications.

            # SpatiaLite case: assume that the string is an hex value
            header = binascii.unhexlify(data[:18]) if isinstance(data, str) else data[:9]
            byte_order, wkb_type, wkb_srid = header[0], header[1:5], header[5:]
            byte_order_marker = "<I" if byte_order else ">I"
            wkb_type_int = (
                int(struct.unpack(byte_order_marker, wkb_type)[0]) if len(wkb_type) == 4 else 0
            )
            if extended is None:
                # Check SRID bit
                extended = False if not wkb_type_int else extended or bool(wkb_type_int & 536870912)
            if extended and srid == -1:
                wkb_srid = struct.unpack(byte_order_marker, wkb_srid)[0]
                srid = int(wkb_srid)
        _SpatialElement.__init__(self, data, srid, extended)

    @staticmethod
    def _wkb_to_hex(data: str | bytes | memoryview) -> str:
        """Convert WKB to hex string."""
        if isinstance(data, str):
            # SpatiaLite case
            return data.lower()
        return str(binascii.hexlify(data), encoding="utf-8").lower()

    @property
    def desc(self) -> str:
        """This element's description string."""
        return self._wkb_to_hex(self.data)

    @staticmethod
    def _data_from_desc(desc) -> bytes:
        desc = desc.encode(encoding="utf-8")
        return binascii.unhexlify(desc)

    def as_wkb(self) -> WKBElement:
        if self.extended:
            if isinstance(self.data, str):
                # SpatiaLite case
                # assume that the string is an hex value
                is_hex = True
                header = binascii.unhexlify(self.data[:10])
                byte_order, wkb_type = header[0], header[1:5]
            else:
                is_hex = False
                byte_order, wkb_type = self.data[0], self.data[1:5]

            byte_order_marker = "<I" if byte_order else ">I"
            wkb_type_int = (
                int(struct.unpack(byte_order_marker, wkb_type)[0]) if len(wkb_type) == 4 else 0
            )
            wkb_type_int &= 3758096383  # Set SRID bit to 0 and keep all other bits

            if is_hex:
                wkb_type_hex = binascii.hexlify(
                    wkb_type_int.to_bytes(4, "little" if byte_order else "big")
                )
                data = self.data[:2] + wkb_type_hex.decode("ascii") + self.data[18:]
            else:
                buffer = bytearray()
                buffer.extend(self.data[:1])
                buffer.extend(struct.pack(byte_order_marker, wkb_type_int))
                buffer.extend(self.data[9:])
                data = memoryview(buffer)
            return WKBElement(data, self.srid, extended=False)
        return WKBElement(self.data, self.srid)

    def as_ewkb(self) -> WKBElement:
        if not self.extended and self.srid != -1:
            if isinstance(self.data, str):
                # SpatiaLite case
                # assume that the string is an hex value
                header = binascii.unhexlify(self.data[:10])
                byte_order, wkb_type = header[0], header[1:5]
            else:
                byte_order, wkb_type = self.data[0], self.data[1:5]
            byte_order_marker = "<I" if byte_order else ">I"
            wkb_type_int = int(
                struct.unpack(byte_order_marker, wkb_type)[0] if len(wkb_type) == 4 else 0
            )
            wkb_type_int |= 536870912  # Set SRID bit to 1 and keep all other bits

            data: str | memoryview
            if isinstance(self.data, str):
                wkb_type_hex = binascii.hexlify(
                    wkb_type_int.to_bytes(4, "little" if byte_order else "big")
                )
                wkb_srid_hex = binascii.hexlify(
                    self.srid.to_bytes(4, "little" if byte_order else "big")
                )
                data = (
                    self.data[:2]
                    + wkb_type_hex.decode("ascii")
                    + wkb_srid_hex.decode("ascii")
                    + self.data[10:]
                )
            else:
                buffer = bytearray()
                buffer.extend(self.data[:1])
                buffer.extend(struct.pack(byte_order_marker, wkb_type_int))
                buffer.extend(struct.pack(byte_order_marker, self.srid))
                buffer.extend(self.data[5:])
                data = memoryview(buffer)

            return WKBElement(data, self.srid, extended=True)
        return WKBElement(self.data, self.srid)


class DynamicWKBElement(WKBElement):
    """This is a subclass of ``WKBElement`` that allows dynamic attributes.

    It is useful when you need to add attributes dynamically to the object.
    """

    __slots__ = ("__dict__", "__weakref__")


class RasterElement(_SpatialElement):
    """Instances of this class wrap a ``raster`` value.

    Raster values read from the database are converted to instances of this type. In
    most cases you won't need to create ``RasterElement`` instances yourself.

    Note::
        This class uses ``__slots__`` to restrict its attributes and improve memory efficiency by
        preventing the creation of a dynamic ``__dict__`` for each instance.
        If you require dynamic attributes or support for weak references, use the
        ``DynamicRasterElement`` subclass, which provides these capabilities.
    """

    __slots__ = ()

    geom_from_extended_version: str = "raster"

    def __init__(self, data: str | bytes | memoryview) -> None:
        # read srid from the WKB (binary or hexadecimal format)
        # The WKB structure is documented in the file
        # raster/doc/RFC2-WellKnownBinaryFormat of the PostGIS sources.
        bin_data: str | bytes | memoryview
        try:
            bin_data = binascii.unhexlify(data[:114])
        except BinasciiError:
            bin_data = data
            data = str(binascii.hexlify(data).decode(encoding="utf-8"))  # type: ignore
        byte_order = bin_data[0]
        srid = bin_data[53:57]
        srid = struct.unpack("<I" if byte_order else ">I", srid)[0]  # type: ignore
        _SpatialElement.__init__(self, data, int(srid), True)

    @property
    def desc(self) -> str:
        """This element's description string."""
        return self.data

    @staticmethod
    def _data_from_desc(desc):
        return desc


class DynamicRasterElement(RasterElement):
    """This is a subclass of ``RasterElement`` that allows dynamic attributes.

    It is useful when you need to add attributes dynamically to the object.
    """

    __slots__ = ("__dict__", "__weakref__")


class CompositeElement(FunctionElement):
    """Instances of this class wrap a Postgres composite type."""

    __slots__ = ("name", "type")

    inherit_cache: bool = False
    """The cache is disabled for this class."""

    def __init__(self, base, field, type_) -> None:
        self.name = field
        self.type = to_instance(type_)

        super().__init__(base)


@compiles(CompositeElement)
def _compile_pgelem(expr, compiler, **kw) -> str:
    return f"({compiler.process(expr.clauses, **kw)}).{expr.name}"


__all__: list[str] = [
    "_SpatialElement",
    "CompositeElement",
    "RasterElement",
    "WKBElement",
    "WKTElement",
    "DynamicRasterElement",
    "DynamicWKBElement",
    "DynamicWKTElement",
]


def __dir__() -> list[str]:
    return __all__
