"""
TiDB TopN Value Decoder

This module provides utilities to decode TiDB TopN values that are encoded as bytes.
TiDB uses a type flag system where each encoded value starts with a byte indicating
the type, followed by the actual encoded data.

Based on TiDB's codec.go implementation:
- intFlag = 3: Signed integers (big-endian 8 bytes)
- uintFlag = 4: Unsigned integers (big-endian 8 bytes)
- floatFlag = 5: Float64 values (IEEE 754 format)
- bytesFlag = 1: Byte strings with length encoding
- compactBytesFlag = 2: Compact byte strings with varint length
- nilFlag = 0: NULL values
"""

import struct
from typing import Any

# Type flags from TiDB codec.go
NIL_FLAG = 0
BYTES_FLAG = 1
COMPACT_BYTES_FLAG = 2
INT_FLAG = 3
UINT_FLAG = 4
FLOAT_FLAG = 5
DECIMAL_FLAG = 6
DURATION_FLAG = 7
VARINT_FLAG = 8
UVARINT_FLAG = 9
JSON_FLAG = 10


def decode_tidb_topn_value(encoded_value: bytes) -> Any:
    """
    Decode a TiDB TopN encoded value back to its original Python type.

    Args:
        encoded_value: The encoded bytes from TiDB TopN value

    Returns:
        The decoded value in appropriate Python type:
        - int for signed integers
        - int for unsigned integers (within Python int range)
        - float for floating point numbers
        - str for string/text values
        - bytes for binary data
        - None for NULL values

    Raises:
        ValueError: If the encoded value format is invalid or unsupported

    Examples:
        >>> decode_tidb_topn_value(b'\\x03\\x80\\x00\\x00\\x00\\x00\\x00\\x00\\x01')
        1
        >>> decode_tidb_topn_value(b'\\x05@\\x08\\x00\\x00\\x00\\x00\\x00\\x00')
        3.0
        >>> decode_tidb_topn_value(b'\\x01hello\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00')
        'hello'
    """
    if not encoded_value or len(encoded_value) == 0:
        raise ValueError("Empty encoded value")

    type_flag = encoded_value[0]
    data = encoded_value[1:]

    if type_flag == NIL_FLAG:
        return None

    elif type_flag == INT_FLAG:
        return _decode_signed_int(data)

    elif type_flag == UINT_FLAG:
        return _decode_unsigned_int(data)

    elif type_flag == FLOAT_FLAG:
        return _decode_float(data)

    elif type_flag == BYTES_FLAG:
        return _decode_bytes(data)

    elif type_flag == COMPACT_BYTES_FLAG:
        return _decode_compact_bytes(data)

    elif type_flag == VARINT_FLAG:
        return _decode_varint(data)

    elif type_flag == UVARINT_FLAG:
        return _decode_uvarint(data)

    elif type_flag == DECIMAL_FLAG:
        # For decimal values, we'll decode as string to preserve precision
        return _decode_decimal_as_string(data)

    elif type_flag == DURATION_FLAG:
        # Duration values are stored as int64 nanoseconds
        duration_ns = _decode_signed_int(data)
        # Convert to seconds as float for compatibility
        return duration_ns / 1_000_000_000.0

    elif type_flag == JSON_FLAG:
        # JSON values are stored as binary JSON format
        # For simplicity, we'll decode as bytes and let the caller handle JSON parsing
        return data

    else:
        raise ValueError(f"Unsupported type flag: {type_flag}")


def _decode_signed_int(data: bytes) -> int:
    """Decode TiDB signed integer (8 bytes big-endian with special encoding)"""
    if len(data) < 8:
        raise ValueError(f"Invalid signed int data length: {len(data)}, expected 8")

    # TiDB uses a special encoding for signed integers to make them sortable
    # We need to reverse the transformation: flip first bit, then interpret as big-endian
    raw_bytes = bytearray(data[:8])
    raw_bytes[0] ^= 0x80  # Flip the sign bit back

    # Convert to signed integer
    value = struct.unpack(">q", raw_bytes)[0]  # '>q' = big-endian signed long long
    return value


def _decode_unsigned_int(data: bytes) -> int:
    """Decode TiDB unsigned integer (8 bytes big-endian)"""
    if len(data) < 8:
        raise ValueError(f"Invalid unsigned int data length: {len(data)}, expected 8")

    # TiDB stores unsigned integers as big-endian 8 bytes
    value = struct.unpack(">Q", data[:8])[0]  # '>Q' = big-endian unsigned long long
    return value


def _decode_float(data: bytes) -> float:
    """Decode TiDB float64 value (8 bytes in IEEE 754 format with special encoding)"""
    if len(data) < 8:
        raise ValueError(f"Invalid float data length: {len(data)}, expected 8")

    # TiDB uses a special encoding for floats to make them sortable
    raw_bytes = bytearray(data[:8])

    # Reverse TiDB's float encoding transformation
    if raw_bytes[0] & 0x80:  # Positive number
        raw_bytes[0] ^= 0x80  # Flip sign bit
    else:  # Negative number
        # Flip all bits
        for i in range(8):
            raw_bytes[i] ^= 0xFF

    # Convert to float
    value = struct.unpack(">d", raw_bytes)[0]  # '>d' = big-endian double
    return value


def _decode_bytes(data: bytes) -> str:
    """Decode TiDB bytes with group encoding"""
    # TiDB encodes bytes in groups of 8 bytes with a marker byte
    # This is a simplified implementation - full implementation would handle group markers

    if not data:
        return ""

    # For now, we'll try to decode as UTF-8 string, handling group markers
    result = []
    i = 0

    while i < len(data):
        # Read 8-byte groups with 1-byte marker
        if i + 8 < len(data):
            group = data[i : i + 8]
            marker = data[i + 8] if i + 8 < len(data) else 0xFF

            # Calculate padding based on marker
            pad_count = 0xFF - marker
            if pad_count > 0:
                # Remove padding bytes
                group = group[:-pad_count]

            result.append(group)
            i += 9  # 8 bytes data + 1 byte marker
        else:
            # Handle remaining bytes
            result.append(data[i:])
            break

    # Combine all groups and decode as UTF-8
    combined = b"".join(result)
    try:
        return combined.decode("utf-8")
    except UnicodeDecodeError:
        # If UTF-8 decoding fails, return as bytes
        return combined.decode("utf-8", errors="replace")


def _decode_compact_bytes(data: bytes) -> str:
    """Decode TiDB compact bytes with varint length prefix"""
    if not data:
        return ""

    # Read varint length
    length, bytes_read = _decode_varint_length(data)
    if bytes_read >= len(data):
        return ""

    # Extract the actual bytes
    actual_data = data[bytes_read : bytes_read + length]

    try:
        return actual_data.decode("utf-8")
    except UnicodeDecodeError:
        return actual_data.decode("utf-8", errors="replace")


def _decode_varint(data: bytes) -> int:
    """Decode variable-length signed integer"""
    if not data:
        raise ValueError("Empty varint data")

    result = 0
    shift = 0

    for byte in data:
        result |= (byte & 0x7F) << shift
        shift += 7

        if (byte & 0x80) == 0:
            break

    # Handle sign bit
    if result & 1:
        return -(result >> 1) - 1
    else:
        return result >> 1


def _decode_uvarint(data: bytes) -> int:
    """Decode variable-length unsigned integer"""
    if not data:
        raise ValueError("Empty uvarint data")

    result = 0
    shift = 0

    for byte in data:
        result |= (byte & 0x7F) << shift
        shift += 7

        if (byte & 0x80) == 0:
            break

    return result


def _decode_varint_length(data: bytes) -> tuple[int, int]:
    """Decode varint length and return (length, bytes_consumed)"""
    if not data:
        return 0, 0

    result = 0
    shift = 0
    bytes_read = 0

    for byte in data:
        result |= (byte & 0x7F) << shift
        shift += 7
        bytes_read += 1

        if (byte & 0x80) == 0:
            break

    return result, bytes_read


def _decode_decimal_as_string(data: bytes) -> str:
    """Decode TiDB decimal value as string to preserve precision"""
    # This is a simplified implementation
    # Full implementation would need to handle TiDB's decimal format

    if not data:
        return "0"

    # For now, try to interpret as a simple string representation
    try:
        return data.decode("utf-8")
    except UnicodeDecodeError:
        # Fallback: return a default decimal string
        return "0.0"


def is_tidb_encoded_value(value: bytes) -> bool:
    """
    Check if a bytes value appears to be TiDB-encoded.

    Args:
        value: Bytes to check

    Returns:
        True if the value appears to be TiDB-encoded, False otherwise
    """
    if not value or len(value) == 0:
        return False

    # Check if first byte is a known type flag
    type_flag = value[0]
    known_flags = {
        NIL_FLAG,
        BYTES_FLAG,
        COMPACT_BYTES_FLAG,
        INT_FLAG,
        UINT_FLAG,
        FLOAT_FLAG,
        DECIMAL_FLAG,
        DURATION_FLAG,
        VARINT_FLAG,
        UVARINT_FLAG,
        JSON_FLAG,
    }

    return type_flag in known_flags


def decode_tidb_topn_value_safe(encoded_value: bytes, default: Any = None) -> Any:
    """
    Safe version of decode_tidb_topn_value that returns a default value on error.

    Args:
        encoded_value: The encoded bytes from TiDB TopN value
        default: Default value to return if decoding fails

    Returns:
        The decoded value or the default value if decoding fails
    """
    try:
        return decode_tidb_topn_value(encoded_value)
    except (ValueError, struct.error, UnicodeDecodeError, IndexError) as e:
        # Log the error in real applications
        print(f"Error decoding TiDB TopN value: {e}")
        return default