import math
from typing import Union


def calculate_checksum(command: Union[int, bytes]) -> bytes:
    # print(command)
    if type(command) == int:
        command = command.to_bytes(64, "big")  # type: ignore
    lrc = 0
    for byte in command:  # type: ignore
        lrc ^= byte
    out = lrc.to_bytes(1, "big")
    # print(out)
    return out


def calculate_long_address(manufacturer_id: int, manufacturer_device_type: int, device_id: bytes):
    out = int.from_bytes(device_id, "big")
    out |= manufacturer_device_type << 24
    out |= manufacturer_id << 32
    return out.to_bytes(5, "big")


def pack_command(address, command_id, data=None):
    # if type(address) == bytes:
    #     address = int.from_bytes(address, "big")
    if type(command_id) == int:
        command_id = command_id.to_bytes(1, "big")
    command = b"\xFF\xFF\xFF\xFF\xFF\xFF"  # preamble
    command += b"\x82"  if len(address) >= 5 else  b"\x02"
    command += address
    command += command_id
    if data is None:
        command += b"\x00"  # byte count
    else:
        # print(len(data), 22222222222)
        command += len(data).to_bytes(1, "big")  # byte count
        command += data  # data
    # print(command[6:])
    command += calculate_checksum(command[6:])
    # print(command)
    return command


PACKED_ASCII_MAP = {
    0x00: '@', 0x01: 'A', 0x02: 'B', 0x03: 'C', 0x04: 'D', 0x05: 'E', 0x06: 'F', 0x07: 'G',
    0x08: 'H', 0x09: 'I', 0x0A: 'J', 0x0B: 'K', 0x0C: 'L', 0x0D: 'M', 0x0E: 'N', 0x0F: 'O',
    0x10: 'P', 0x11: 'Q', 0x12: 'R', 0x13: 'S', 0x14: 'T', 0x15: 'U', 0x16: 'V', 0x17: 'W',
    0x18: 'X', 0x19: 'Y', 0x1A: 'Z', 0x1B: '[', 0x1C: '\\', 0x1D: ']', 0x1E: '^', 0x1F: '_',  # 这里应该是 '_' 而不是 '-'
    0x20: ' ', 0x21: '!', 0x22: '"', 0x23: '=', 0x24: '$', 0x25: '%', 0x26: '&', 0x27: '`',  # 注意这里是反引号 '`'
    0x28: '(', 0x29: ')', 0x2A: '*', 0x2B: '+', 0x2C: ',', 0x2D: '-', 0x2E: '.', 0x2F: '/',
    0x30: '0', 0x31: '1', 0x32: '2', 0x33: '3', 0x34: '4', 0x35: '5', 0x36: '6', 0x37: '7',
    0x38: '8', 0x39: '9', 0x3A: ':', 0x3B: ';', 0x3C: '<', 0x3D: '=', 0x3E: '>', 0x3F: '?'
}



def unpack_packed_ascii(data: bytes, expected_len: int) -> str:
    """Unpacks a byte array into a string using the HART Packed ASCII 6-bit encoding by processing in chunks."""
    chars = []
    # Process data in 3-byte chunks, as 4 characters (24 bits) fit into 3 bytes (24 bits)
    for i in range(0, len(data), 3):
        chunk = data[i:i+3]
        
        # Pad chunk to 3 bytes if it's a partial chunk at the end
        if len(chunk) < 3:
            chunk += b'\x00' * (3 - len(chunk))
        
        byte1, byte2, byte3 = chunk
        
        # Unpack 4 6-bit characters from the 3 8-bit bytes
        c1 = byte1 >> 2
        c2 = ((byte1 & 0x03) << 4) | (byte2 >> 4)
        c3 = ((byte2 & 0x0F) << 2) | (byte3 >> 6)
        c4 = byte3 & 0x3F
        
        codes = [c1, c2, c3, c4]
        for code in codes:
            chars.append(PACKED_ASCII_MAP.get(code, '?'))

    # Join and then trim to the exact expected length. Do not rstrip() as spaces can be valid.
    return "".join(chars[:expected_len])




# Build a reverse map for packing, handling duplicates deterministically
REVERSE_PACKED_ASCII_MAP = {}
# Iterate through the original map, sorted by code value (key)
for code, char in sorted(PACKED_ASCII_MAP.items(), key=lambda item: item[0]):
    # This ensures that for characters with multiple codes, the one with the highest code value is chosen.
    REVERSE_PACKED_ASCII_MAP[char] = code


def pack_packed_ascii(text: Union[str, bytes], expected_len: int) -> bytes:
    """
    Packs a string or bytes into a HART Packed ASCII byte array.
    This function is the symmetrical inverse of the unpack_packed_ascii function.
    """
    # Defensively decode if bytes are passed in
    if isinstance(text, bytes):
        try:
            # Try decoding with ascii, fall back to latin-1 which never fails
            text = text.decode('ascii')
        except UnicodeDecodeError:
            text = text.decode('latin-1')

    # 1. Pad/truncate the input string to the exact expected length
    padded_text = text.ljust(expected_len, ' ')
    padded_text = padded_text[:expected_len]

    # 2. Convert all characters to their corresponding 6-bit codes
    codes = [REVERSE_PACKED_ASCII_MAP.get(c, 0x3F) for c in padded_text] # Default to '?' (0x3F)
    
    packed_bytes = bytearray()
    
    # 3. Process the codes in chunks of 4
    for i in range(0, len(codes), 4):
        # Get a chunk of up to 4 codes.
        chunk = codes[i:i+4]
        
        # If it's a partial chunk at the end, pad with space codes to make a full chunk of 4.
        while len(chunk) < 4:
            chunk.append(REVERSE_PACKED_ASCII_MAP[' '])
        
        c1, c2, c3, c4 = chunk
        
        # 4. Pack the 4 6-bit codes into 3 8-bit bytes, mirroring the unpacking logic
        byte1 = (c1 << 2) | (c2 >> 4)
        byte2 = ((c2 & 0x0F) << 4) | (c3 >> 2)
        byte3 = ((c3 & 0x03) << 6) | c4
        
        packed_bytes.extend([byte1, byte2, byte3])
        
    # 5. Calculate the exact number of bytes required for the original expected length
    num_bytes = math.ceil(expected_len * 6 / 8)
    
    # 6. Get the core byte array, trimmed to the precise required size
    result = bytes(packed_bytes[:num_bytes])

    # For Command 17 (Write Message), the spec requires a fixed 24-byte data field.
    # This corresponds to an expected character length of 32.
    if expected_len == 32:
        # Pad the result to exactly 24 bytes.
        return result.ljust(24, b'\x00')
    
    return result