import base64
from enum import Enum
from queryservice_client.protobuf.rdbms.query_pb2 import RowData, Null, Endianness

from datetime import datetime, timezone, timedelta
from typing import Any, List
from decimal import Decimal

class Type(Enum):
    NULL = 0
    DOUBLE = 1
    FLOAT = 2
    INT32 = 3
    INT64 = 4
    UINT32 = 5
    UINT64 = 6
    BOOL = 7
    STRING = 8
    BYTES = 9
    DATE = 10
    ARRAY = 11
    STRUCT = 12
    BIGINT = 13
    BIGDECIMAL = 14

class Data:
    def __init__(self, data: RowData):
        self.data = data

    @staticmethod
    def create(row_data: RowData) -> 'Data':
        data_case = row_data.WhichOneof('data')
        if data_case == 'null':
            return NullData(row_data)
        elif data_case == 'double':
            return DoubleData(row_data)
        elif data_case == 'float':
            return FloatData(row_data)
        elif data_case == 'int32':
            return Int32Data(row_data)
        elif data_case == 'int64':
            return Int64Data(row_data)
        elif data_case == 'uint32':
            return Uint32Data(row_data)
        elif data_case == 'uint64':
            return Uint64Data(row_data)
        elif data_case == 'bool':
            return BoolData(row_data)
        elif data_case == 'string':
            return StringData(row_data)
        elif data_case == 'bytes':
            return BytesData(row_data)
        elif data_case == 'date':
            return DateData(row_data)
        elif data_case == 'array':
            return ArrayData(row_data)
        elif data_case == 'struct':
            return StructData(row_data)
        elif data_case == 'bigInt':
            return BigIntData(row_data)
        elif data_case == 'bigDecimal':
            return BigDecimalData(row_data)
        else:
            raise RuntimeError(f"Unknown datatype: {data_case}")

    def type(self) -> Type:
        raise NotImplementedError

    def get(self) -> Any:
        raise NotImplementedError

    def __eq__(self, other: Any) -> bool:
        if isinstance(other, Data):
            return self.data == other.data
        return False

    def __hash__(self) -> int:
        return hash(self.data)

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


class NullData(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.NULL

    def get(self) -> None:
        return None

    def __str__(self) -> str:
        return "NULL"
    
class Int64Data(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.INT64

    def get(self) -> int:
        return self.data.int64
    
class Int32Data(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.INT32

    def get(self) -> int:
        return self.data.int32
    
class FloatData(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.FLOAT

    def get(self) -> float:
        return self.data.float

class DoubleData(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.DOUBLE

    def get(self) -> float:
        return self.data.double

    def __str__(self) -> str:
        return str(self.get())
    
class DateData(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.DATE

    def get(self) -> datetime:
        date = self.data.date
        epoch_seconds = date.timestamp.seconds
        nanos = date.timestamp.nanos
        offset_minutes = date.timezoneOffsetMinutes
        offset = timedelta(minutes=offset_minutes)
        return datetime.fromtimestamp(epoch_seconds + nanos / 1e9, tz=timezone(offset))

    def __str__(self) -> str:
        return self.get().isoformat()

class StringData(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.STRING

    def get(self) -> str:
        return self.data.string

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


class BytesData(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.BYTES

    def get(self) -> str:
        return base64.b64encode(self.data.bytes).decode('utf-8')
    

class BoolData(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.BOOL

    def get(self) -> bool:
        return self.data.bool

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

class BigIntData(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.BIGINT

    def get(self) -> int:
        big_int = self.data.bigInt
        bytes_value = big_int.value
        if big_int.endianness == Endianness.LITTLE:
            bytes_value = bytes_value[::-1]
        return int.from_bytes(bytes_value, byteorder='big', signed=True)
    

class BigDecimalData(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.BIGDECIMAL

    def get(self) -> Decimal:
        big_decimal = self.data.bigDecimal
        scale = big_decimal.scale
        big_int = big_decimal.unscaledValue
        bytes_value = big_int.value
        if big_int.endianness == Endianness.LITTLE:
            bytes_value = bytes_value[::-1]
        unscaled_value = int.from_bytes(bytes_value, byteorder='big', signed=True)
        return Decimal(unscaled_value).scaleb(-scale)

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

class ArrayData(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.ARRAY

    def get(self) -> List[Data]:
        return [Data.create(element) for element in self.data.array.array]

    def __str__(self) -> str:
        values = ", ".join(f'"{data}"' if data.type() == Type.STRING else str(data) for data in self.get())
        return f"[{values}]"
    

class StructData(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.STRUCT

    def get(self) -> List[Data]:
        return [Data.create(member) for member in self.data.struct.members]

    def __str__(self) -> str:
        values = ", ".join(f'"{data}"' if data.type() == Type.STRING else str(data) for data in self.get())
        return f"({values})"
    

class Uint32Data(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.UINT32

    def get(self) -> int:
        return self.data.uint32
    

class Uint64Data(Data):
    def __init__(self, data: RowData):
        super().__init__(data)

    def type(self) -> Type:
        return Type.UINT64

    def get(self) -> int:
        return self.data.uint64