# Copyright 2026, Jiri Eischmann <jiri@eischmann.cz> and the meshy contributors
# SPDX-License-Identifier: GPL-3.0-or-later

"""Lightweight MeshCore packet decoding and channel crypto.

Replaces meshcoredecoder for the subset of features Meshy needs:
- Channel hash calculation (SHA-256)
- GroupText payload decryption (HMAC-SHA256 + AES-128-ECB)
- Minimal packet header/path parsing for LOG_DATA frames
"""

import hashlib
import hmac
import struct
from typing import Optional


# Payload type constants (from header bits 2-5)
PAYLOAD_ADVERT = 0x04
PAYLOAD_GROUP_TEXT = 0x05


def channel_hash(psk_hex: str) -> str:
    """Calculate the 1-byte channel hash from a PSK hex string.

    This is SHA-256(psk_bytes)[0] as a 2-char lowercase hex string.
    """
    h = hashlib.sha256(bytes.fromhex(psk_hex)).digest()
    return f'{h[0]:02x}'


def decrypt_group_text(ciphertext: bytes, psk: bytes) -> Optional[dict]:
    """Decrypt a GroupText payload using the channel PSK.

    Args:
        ciphertext: The encrypted payload (after channel_hash byte).
                    First 2 bytes are HMAC-SHA256 MAC, rest is AES-128-ECB ciphertext.
        psk: The 16-byte channel pre-shared key.

    Returns:
        Dict with 'sender', 'message', 'timestamp', 'flags' or None on failure.
    """
    if len(ciphertext) < 3:
        return None

    mac = ciphertext[:2]
    encrypted = ciphertext[2:]

    # Verify HMAC-SHA256: key is PSK padded to 32 bytes with zeros
    key32 = psk + b'\x00' * (32 - len(psk)) if len(psk) < 32 else psk[:32]
    expected_mac = hmac.new(key32, encrypted, hashlib.sha256).digest()
    if mac[0] != expected_mac[0] or mac[1] != expected_mac[1]:
        return None

    # Decrypt with AES-128-ECB
    try:
        from Crypto.Cipher import AES
        key16 = psk[:16] if len(psk) >= 16 else psk + b'\x00' * (16 - len(psk))
        cipher = AES.new(key16, AES.MODE_ECB)
        plaintext = cipher.decrypt(encrypted)
    except ImportError:
        return None

    # Parse plaintext: timestamp(4 LE) + flags(1) + "sender: message"
    if len(plaintext) < 6:
        return None

    timestamp = struct.unpack('<I', plaintext[:4])[0]
    flags = plaintext[4]
    text = plaintext[5:].decode('utf-8', 'ignore').rstrip('\x00')

    # Split "sender: message" on first ": "
    if ': ' in text:
        sender, message = text.split(': ', 1)
    else:
        sender = ''
        message = text

    return {
        'sender': sender,
        'message': message,
        'timestamp': timestamp,
        'flags': flags,
    }


def decode_log_packet(raw: bytes) -> Optional[dict]:
    """Parse a raw mesh packet from LOG_DATA, extract GroupText if present.

    Parses the MeshCore packet header, skips path data, and extracts
    the payload. Only GroupText payloads are fully decoded.

    Args:
        raw: Raw packet bytes (after LOG_DATA framing).

    Returns:
        Dict with 'payload_type', 'path' (list of hex hop strings),
        'channel_hash', and 'decrypted' (if GroupText), or None.
    """
    if len(raw) < 2:
        return None

    header = raw[0]
    route_type = header & 0x03
    payload_type = (header >> 2) & 0x0F

    offset = 1

    # Parse transport codes for TC_FLOOD (0) and TC_DIRECT (3)
    transport_from = ''
    transport_to = ''
    if route_type == 0 or route_type == 3:
        if len(raw) < offset + 4:
            return None
        transport_from = raw[offset:offset + 2].hex()
        transport_to = raw[offset + 2:offset + 4].hex()
        offset += 4

    # Parse path length byte
    if len(raw) < offset + 1:
        return None
    path_len_byte = raw[offset]
    offset += 1

    # Decode path: bits 7:6 = hash_size selector, bits 5:0 = hop count
    hash_size = (path_len_byte >> 6) + 1
    hop_count = path_len_byte & 63
    path_byte_length = hop_count * hash_size

    # Legacy fallback: reserved hash_size (4) or path extends past packet
    if hash_size == 4 or offset + path_byte_length > len(raw):
        path_byte_length = min(path_len_byte, len(raw) - offset, 64)
        hop_count = path_byte_length
        hash_size = 1

    # Extract path hops as hex strings
    path = []
    for i in range(hop_count):
        hop_start = offset + i * hash_size
        hop_bytes = raw[hop_start:hop_start + hash_size]
        path.append(hop_bytes.hex())
    offset += path_byte_length

    # Extract payload
    payload = raw[offset:]

    pkt_hash = hashlib.sha256(bytes([payload_type]) + payload).hexdigest()[:16]

    result = {
        'payload_type': payload_type,
        'route_type': route_type,
        'path': path,
        'path_hash_size': hash_size,
        'pkt_hash': pkt_hash,
        'transport_from': transport_from,
        'transport_to': transport_to,
    }

    if payload_type == PAYLOAD_ADVERT:
        advert = parse_advert_payload(payload)
        if advert:
            result['advert'] = advert

    # Decode GroupText payload: channel_hash(1) + mac(2) + ciphertext
    elif payload_type == PAYLOAD_GROUP_TEXT and len(payload) >= 3:
        result['channel_hash'] = f'{payload[0]:02x}'
        result['ciphertext'] = payload[1:]  # mac + encrypted data

    return result


def parse_advert_payload(payload: bytes) -> Optional[dict]:
    """Parse an ADVERT payload from a LOG_DATA frame.

    Format: pubkey(32) + timestamp(4) + signature(64) + flags(1)
            + [lat(4) + lon(4) if flag 0x10] + [name(null-terminated) if flag 0x80]
    """
    if len(payload) < 101:  # 32 + 4 + 64 + 1
        return None
    pub_key = payload[:32].hex()
    timestamp = struct.unpack_from('<I', payload, 32)[0]
    # Skip 64-byte signature
    flags = payload[100]
    node_type = flags & 0x0F
    has_location = bool(flags & 0x10)
    has_name = bool(flags & 0x80)

    offset = 101
    lat, lon = None, None
    if has_location and len(payload) >= offset + 8:
        lat = struct.unpack_from('<i', payload, offset)[0] / 1e6
        lon = struct.unpack_from('<i', payload, offset + 4)[0] / 1e6
        if abs(lat) < 1e-6 and abs(lon) < 1e-6:
            lat, lon = None, None
        offset += 8

    name = ''
    if has_name and len(payload) > offset:
        end = payload.index(0, offset) if 0 in payload[offset:] else len(payload)
        name = payload[offset:end].decode('utf-8', 'ignore')

    return {
        'public_key': pub_key,
        'type': node_type,
        'adv_name': name,
        'last_advert': timestamp,
        'adv_lat': lat,
        'adv_lon': lon,
    }
