from __future__ import annotations

"""
Extended Types
"""

import contextlib
from enum import Enum
from sqlalchemy import types
from sqlalchemy.schema import Sequence
from sqlalchemy.engine.interfaces import Dialect
from sqlalchemy.sql import operators
from sqlalchemy.dialects.postgresql import JSONB, ARRAY, JSON
from sqlalchemy.dialects.postgresql.base import ischema_names

from lzl import load
import typing as t

# if load.t.TYPE_CHECKING:
#     import fileio
# else:
#     fileio = load.LazyLoad("fileio", install_missing=True, install_options={'package': 'file-io'})

# if t.TYPE_CHECKING:
#     from fileio import FileLike

from lzl.io import File, FileLike

class FileField(types.TypeDecorator):

    impl = types.String
    cache_ok = True
    mysql_default_length = 255

    def coerce_compared_value(self, op, value):
        """
        Coerce the value to be compared.
        """
        return types.String(length = self.mysql_default_length) if op in (operators.like_op, operators.not_like_op) else self
    
    def process_bind_param(self, value: t.Optional[t.Union[str, 'FileLike']], dialect):
        """
        Return the value to be stored in the database.
        """
        if value is None: return None
        if isinstance(value, str): return value
        return value.as_posix() \
            if hasattr(value, 'as_posix') \
                else str(value)

    def process_result_value(self, value: t.Optional[t.Union[str, t.Any]], dialect):
        """
        Process the value retrieved from the database.
        """
        return File(value) if value is not None else None


class JsonString(types.TypeDecorator):
    """
    A String type that stores JSON data.
    """
    impl = types.String
    cache_ok = True
    hashable = True

    def __init__(
        self,
        *args,
        **kwargs,
    ) -> None:
        """
        Initializes the JsonString t.Type
        """
        super().__init__(*args, **kwargs)
        from lzl.io.ser import get_serializer
        self.ser = get_serializer('json')
        self.ser.enforce_string_value = True

    def load_dialect_impl(self, dialect):
        """
        Load the dialect implementation.
        """
        return dialect.type_descriptor(self.impl)

    def process_bind_param(self, value, dialect) -> t.Optional[str]:
        """
        Process the value to be stored in the database.
        """
        if value is not None:
            value = self.ser.dumps(value)
        return value

    def process_result_value(self, value, dialect) ->  t.Optional[t.Union[t.Dict[str, t.Any], t.List[t.Dict[str, t.Any]]]]:
        """
        Process the value retrieved from the database.
        """
        if value is not None:
            value = self.ser.loads(value)
        return value


class SerializedBinary(types.TypeDecorator):
    """
    A LargeBinary type that stores serialized data.
    """
    impl = types.LargeBinary
    cache_ok = True
    hashable = True

    def __init__(
        self,
        *args,
        serializer: t.Optional[str] = 'json',
        compression: t.Optional[str] = 'zstd',
        compression_level: int = 3,
        validate_raw: t.Optional[bool] = True,
        **kwargs,
    ) -> None:
        """
        Initializes the SerializedBinary t.Type

        - Supports Serialization: json, pickle, msgpack
        - Supports Compression: zstd, lz4, zlib, and bz2
        """
        super().__init__(*args, **kwargs)
        from lzl.io.ser import get_serializer
        self.ser = get_serializer(serializer, compression = compression, compression_level = compression_level)
        self.validate_raw = validate_raw

    
    def load_dialect_impl(self, dialect):
        """
        Load the dialect implementation.
        """
        return dialect.type_descriptor(self.impl)

    def process_bind_param(self, value: t.Union[str, bytes, t.Any], dialect) -> t.Optional[bytes]:
        """
        Process the value to be stored in the database.
        """
        if value is not None:
            if self.validate_raw:
                if isinstance(value, bytes): return value
                if isinstance(value, str): return value.encode()
            value = self.ser.dumps(value)
        return value

    def process_result_value(self, value, dialect) ->  t.Optional[t.Union[t.Dict[str, t.Any], t.List[t.Dict[str, t.Any]]]]:
        """
        Process the value retrieved from the database.
        """
        if value is not None: value = self.ser.loads(value)
        return value



DateTimeUTC: types.DateTime = types.DateTime(timezone = True)

IntList: ARRAY = ARRAY(types.Integer, dimensions = 1)
BigIntList: ARRAY = ARRAY(types.BigInteger, dimensions = 1)
StringList: ARRAY = ARRAY(types.Text, dimensions = 1)
DateTimeUTCList: ARRAY = ARRAY(DateTimeUTC, dimensions = 1)
# AutoSequence: types.BigInteger = 

with contextlib.suppress(ImportError):
    from sqlalchemy_json import NestedMutableJson, MutableDict, mutable_json_type

    JsonStringField: t.Type[MutableDict] = mutable_json_type(dbtype = JsonString, nested=True)
    JsonStringNestedField: t.Type[NestedMutableJson] = mutable_json_type(dbtype = JsonString, nested = True) # type: ignore


    JSONBField: t.Type[MutableDict] = mutable_json_type(dbtype = JSONB)
    JSONBNestedField: t.Type[NestedMutableJson] = mutable_json_type(dbtype = JSONB, nested = True) # type: ignore


# for reflection
ischema_names['file'] = FileField
ischema_names['jsonstring'] = JsonString
ischema_names['serializedbinary'] = SerializedBinary

# Try to import pgvector
try:
    from pgvector.sqlalchemy import Vector

except ImportError:
    import numpy as np

    def from_db(value):
        # could be ndarray if already cast by lower-level driver
        if value is None or isinstance(value, np.ndarray):
            return value
        return np.array(value[1:-1].split(','), dtype=np.float32)
        
    def to_db(value, dim=None):
        if value is None: return value
        if isinstance(value, np.ndarray):
            if value.ndim != 1:
                raise ValueError('expected ndim to be 1')
            if not np.issubdtype(value.dtype, np.integer) and not np.issubdtype(value.dtype, np.floating):
                raise ValueError('dtype must be numeric')
            value = value.tolist()
        if dim is not None and len(value) != dim:
            raise ValueError('expected %d dimensions, not %d' % (dim, len(value)))
        return '[' + ','.join([str(float(v)) for v in value]) + ']'


    class Vector(types.UserDefinedType):
        cache_ok = True
        _string = types.String()

        def __init__(self, dim=None):
            super(types.UserDefinedType, self).__init__()
            self.dim = dim

        def get_col_spec(self, **kw):
            return "VECTOR" if self.dim is None else "VECTOR(%d)" % self.dim

        def bind_processor(self, dialect):
            def process(value):
                return to_db(value, self.dim)
            return process

        def literal_processor(self, dialect):
            string_literal_processor = self._string._cached_literal_processor(dialect)
            def process(value):
                return string_literal_processor(to_db(value, self.dim))
            return process

        def result_processor(self, dialect, coltype):
            def process(value):
                return from_db(value)
            return process

        class comparator_factory(types.UserDefinedType.Comparator):
            def l2_distance(self, other):
                return self.op('<->', return_type=types.Float)(other)

            def max_inner_product(self, other):
                return self.op('<#>', return_type=types.Float)(other)

            def cosine_distance(self, other):
                return self.op('<=>', return_type=types.Float)(other)

    ischema_names['vector'] = Vector

# from pgvector.asyncpg import Vector

# https://github.com/pgvector/pgvector


class VectorDistance(str, Enum):
    """
    Supported distance functions are:

    <-> - L2 distance
    <#> - (negative) inner product
    <=> - cosine distance
    <+> - L1 distance (added in 0.7.0)
    <~> - Hamming distance (binary vectors, added in 0.7.0)
    <%> - Jaccard distance (binary vectors, added in 0.7.0)
    """
    L2_DISTANCE = '<->'
    INNER_PRODUCT = '<#>'
    COSINE_DISTANCE = '<=>'
    L1_DISTANCE = '<+>'
    HAMMING_DISTANCE = '<~>'
    JACCARD_DISTANCE = '<%>'

    @property
    def vector_op(self) -> str:
        """
        Returns the Vector Op Name

        vector_l2_ops = L2 distance
        vector_ip_ops = (Negative) Inner Product
        vector_cosine_ops = Cosine Distance
        vector_l1_ops = L1 distance
        vector_hamming_ops = Hamming distance
        vector_jaccard_ops = Jaccard distance
        """
        if self == VectorDistance.L2_DISTANCE:
            return 'vector_l2_ops'
        elif self == VectorDistance.INNER_PRODUCT:
            return 'vector_ip_ops'
        elif self == VectorDistance.COSINE_DISTANCE:
            return 'vector_cosine_ops'
        elif self == VectorDistance.L1_DISTANCE:
            return 'vector_l1_ops'
        elif self == VectorDistance.HAMMING_DISTANCE:
            return 'vector_hamming_ops'
        elif self == VectorDistance.JACCARD_DISTANCE:
            return 'vector_jaccard_ops'
        return 'vector_ip_ops'

    
    @classmethod
    def parse(cls, value: str) -> 'VectorDistance':
        """
        Parses the Vector Distance
        """
        value = value.lower()
        if value in {'euclidean', 'l2_distance', 'l2', 'vector_l2_ops', '<->'}:
            return cls.L2_DISTANCE
        if value in {'inner_product', 'ip', 'negative_inner_product', 'max_inner_product', 'vector_ip_ops', '<#>'}:
            return cls.INNER_PRODUCT
        if value in {'cosine_distance', 'cosine', 'cosine_similarity', 'vector_cosine_ops', '<=>'}:
            return cls.COSINE_DISTANCE
        if value in {'l1', 'vector_l1_ops', '<+>'}:
            return cls.L1_DISTANCE
        if value in {'hamming', 'vector_hamming_ops', '<~>'}:
            return cls.HAMMING_DISTANCE
        if value in {'jaccard', 'vector_jaccard_ops', '<%>'}:
            return cls.JACCARD_DISTANCE
        raise ValueError(f"Invalid/Unsupported Vector Distance: {value}")
