"""
SMP serial client for MCUboot serial recovery (in-process, no external tools).

Implements the mcumgr-over-UART transport used by MCUboot's boot_serial:
  - frame:  [0x06 0x09 | base64(payload) | 0x0a]   (first frame)
           [0x04 0x14 | base64(payload) | 0x0a]   (continuation frames)
  - payload: [len:2 BE][SMP packet:8B hdr + CBOR][crc16-ccitt:2 BE]
  - CRC16-CCITT (poly 0x1021, init 0), appended BIG-endian (matches MCUboot,
    which uses htons + crc16_ccitt(0, ...) and checks for a zero result).

MCUboot serial recovery always writes uploaded images to slot 0 (primary),
so no test/confirm marking is needed - upload then reset is the whole flow.
"""
import struct
import base64
import hashlib
import time
import sys
import serial

# mcumgr over UART framing constants (match MCUboot boot_serial)
SMP_HEADER = b"\x06\x09"
SMP_FRAG_HEADER = b"\x04\x14"
SMP_NEWLINE = b"\x0a"
FRAME_MTU = 124  # MCUboot BOOT_SERIAL_FRAME_MTU: 127 - 2 - 1


class SMPError(Exception):
    """SMP communication error."""


def crc16_ccitt(data: bytes, crc: int = 0) -> int:
    """CRC16-CCITT (poly 0x1021, init 0), same as Zephyr crc16_ccitt."""
    for byte in data:
        crc ^= byte << 8
        for _ in range(8):
            if crc & 0x8000:
                crc = ((crc << 1) ^ 0x1021) & 0xFFFF
            else:
                crc = (crc << 1) & 0xFFFF
    return crc


# ── CBOR encoding (requests) ───────────────────────────────────
def cbor_uint(val: int) -> bytes:
    if val < 24:
        return bytes([val])
    elif val < 256:
        return bytes([0x18, val])
    elif val < 65536:
        return bytes([0x19, val >> 8, val & 0xFF])
    else:
        return bytes([0x1a, (val >> 24) & 0xFF, (val >> 16) & 0xFF,
                      (val >> 8) & 0xFF, val & 0xFF])


def cbor_bytes(data: bytes) -> bytes:
    n = len(data)
    if n < 24:
        return bytes([0x40 + n]) + data
    elif n < 256:
        return bytes([0x58, n]) + data
    else:
        return bytes([0x59, n >> 8, n & 0xFF]) + data


def cbor_str(s: str) -> bytes:
    b = s.encode('utf-8')
    n = len(b)
    if n < 24:
        return bytes([0x60 + n]) + b
    elif n < 256:
        return bytes([0x78, n]) + b
    else:
        return bytes([0x79, n >> 8, n & 0xFF]) + b


# ── CBOR decoding (responses) ──────────────────────────────────
def _cbor_decode(data: bytes, off: int = 0):
    """Decode one CBOR item. Returns (value, new_offset).

    Supports unsigned ints, byte/str strings, maps and arrays - enough
    for MCUboot boot_serial responses ({rc, off} maps and image lists).
    """
    if off >= len(data):
        raise SMPError("CBOR: truncated data")
    b = data[off]
    off += 1

    if b <= 0x17:
        return b, off
    if b == 0x18:
        return data[off], off + 1
    if b == 0x19:
        return struct.unpack('>H', data[off:off + 2])[0], off + 2
    if b == 0x1a:
        return struct.unpack('>I', data[off:off + 4])[0], off + 4
    if 0x40 <= b <= 0x5b:  # byte string
        if b <= 0x57:
            n = b - 0x40
        elif b == 0x58:
            n = data[off]; off += 1
        elif b == 0x59:
            n = struct.unpack('>H', data[off:off + 2])[0]; off += 2
        else:
            n = struct.unpack('>I', data[off:off + 4])[0]; off += 4
        return data[off:off + n], off + n
    if 0x60 <= b <= 0x7b:  # text string
        if b <= 0x77:
            n = b - 0x60
        elif b == 0x78:
            n = data[off]; off += 1
        elif b == 0x79:
            n = struct.unpack('>H', data[off:off + 2])[0]; off += 2
        else:
            n = struct.unpack('>I', data[off:off + 4])[0]; off += 4
        return data[off:off + n].decode('utf-8', 'replace'), off + n
    if 0x80 <= b <= 0x9b:  # array
        if b <= 0x97:
            n = b - 0x80
        elif b == 0x98:
            n = data[off]; off += 1
        elif b == 0x99:
            n = struct.unpack('>H', data[off:off + 2])[0]; off += 2
        else:
            n = struct.unpack('>I', data[off:off + 4])[0]; off += 4
        out = []
        for _ in range(n):
            v, off = _cbor_decode(data, off)
            out.append(v)
        return out, off
    if 0xa0 <= b <= 0xbb:  # map
        if b <= 0xb7:
            n = b - 0xa0
        elif b == 0xb8:
            n = data[off]; off += 1
        elif b == 0xb9:
            n = struct.unpack('>H', data[off:off + 2])[0]; off += 2
        else:
            n = struct.unpack('>I', data[off:off + 4])[0]; off += 4
        out = {}
        for _ in range(n):
            k, off = _cbor_decode(data, off)
            v, off = _cbor_decode(data, off)
            out[k] = v
        return out, off
    raise SMPError(f"CBOR: unsupported type byte 0x{b:02x}")


class SMPClient:
    """Minimal, verified SMP serial client for MCUboot serial recovery."""

    OP_READ = 0
    OP_READ_RSP = 1
    OP_WRITE = 2
    OP_WRITE_RSP = 3

    GROUP_OS = 0
    GROUP_IMAGE = 1

    IMAGE_STATE_READ = 0
    IMAGE_UPLOAD = 1
    IMAGE_STATE_WRITE = 2

    OS_RESET = 5

    def __init__(self, port: str, baudrate: int = 115200, log=None):
        self.port = port
        self.baudrate = baudrate
        self.ser: serial.Serial | None = None
        self._log = log or (lambda m: print(f"[SMP] {m}", file=sys.stderr))

    # ── connection ──────────────────────────────────────────
    def open(self):
        self.ser = serial.Serial(self.port, self.baudrate, timeout=0.1)
        # 清掉设备刚重启时可能残留的字节
        time.sleep(0.2)
        self.ser.reset_input_buffer()
        self._log(f"opened {self.port} @ {self.baudrate}")

    def close(self):
        if self.ser:
            try:
                self.ser.close()
            except Exception:
                pass
            self.ser = None

    # ── packet / frame building ─────────────────────────────
    def _build_smp_packet(self, op: int, group: int, cmd: int,
                          payload: bytes = b"", seq: int = 0) -> bytes:
        """SMP packet: 8B header + CBOR payload + 2B BE CRC, wrapped in
        [len:2 BE] transport prefix."""
        hdr = bytearray(8)
        hdr[0] = op
        hdr[1] = 0  # flags
        hdr[2:4] = struct.pack('>H', len(payload) + 8)
        hdr[4:6] = struct.pack('>H', group)
        hdr[6] = seq
        hdr[7] = cmd
        packet = bytes(hdr) + payload
        crc = crc16_ccitt(packet)
        return struct.pack('>H', len(packet) + 2) + packet + struct.pack('>H', crc)

    def _send_frame(self, data: bytes):
        """Base64-encode and send as mcumgr UART frame(s) with fragmentation."""
        encoded = base64.b64encode(data)
        offset = 0
        first = True
        while offset < len(encoded):
            chunk = encoded[offset:offset + FRAME_MTU]
            header = SMP_HEADER if first else SMP_FRAG_HEADER
            first = False
            self.ser.write(header + chunk + SMP_NEWLINE)
            offset += len(chunk)
        self.ser.flush()

    # ── response receiving ──────────────────────────────────
    def _recv_packet(self, timeout: float = 10.0) -> bytes:
        """Read one complete decoded SMP packet (transport payload without
        the 2-byte length prefix / 2-byte CRC)."""
        start = time.time()
        buf = bytearray()
        while time.time() - start < timeout:
            if self.ser and self.ser.in_waiting:
                buf.extend(self.ser.read(self.ser.in_waiting))
                while b'\n' in buf:
                    nl = buf.index(b'\n')
                    line = bytes(buf[:nl])
                    del buf[:nl + 1]
                    if not line:
                        continue
                    if line[:2] not in (SMP_HEADER, SMP_FRAG_HEADER):
                        self._log(f"skip non-SMP line: {line[:16].hex()}")
                        continue
                    encoded = line[2:]
                    try:
                        decoded = base64.b64decode(encoded, validate=True)
                    except Exception as e:
                        self._log(f"base64 decode failed: {e}")
                        continue
                    if len(decoded) < 12:  # 2 len + 8 hdr + 2 crc
                        self._log("decoded too short")
                        continue
                    total_len = struct.unpack('>H', decoded[:2])[0]
                    if len(decoded) < 2 + total_len:
                        self._log(f"short packet: {len(decoded)} < {2 + total_len}")
                        continue
                    packet = decoded[2:2 + total_len - 2]
                    wire_crc = struct.unpack('>H', decoded[2 + total_len - 2:2 + total_len])[0]
                    if crc16_ccitt(packet) != wire_crc:
                        self._log("CRC mismatch")
                        continue
                    return packet
            time.sleep(0.01)
        raise SMPError(f"timeout waiting for SMP response ({timeout}s)")

    def _request(self, op: int, group: int, cmd: int, payload: bytes = b"",
                 seq: int = 0, timeout: float = 10.0):
        """Send a request and return the response header + payload."""
        self._send_frame(self._build_smp_packet(op, group, cmd, payload, seq))
        resp = self._recv_packet(timeout)
        if len(resp) < 8:
            raise SMPError("response too short")
        r_op = resp[0]
        r_group = struct.unpack('>H', resp[4:6])[0]
        r_seq = resp[6]
        r_cmd = resp[7]
        r_payload = resp[8:]
        if r_op != (self.OP_READ_RSP if op == self.OP_READ else self.OP_WRITE_RSP):
            self._log(f"unexpected op {r_op} for request op {op}")
        return r_group, r_op, r_seq, r_cmd, r_payload

    @staticmethod
    def _rc(payload: bytes) -> int:
        """Extract 'rc' from a CBOR map response; missing means 0 (OK)."""
        if not payload:
            return 0
        try:
            val, _ = _cbor_decode(payload)
            if isinstance(val, dict):
                return int(val.get('rc', 0))
        except SMPError:
            pass
        return 0

    @staticmethod
    def _off(payload: bytes):
        """Extract 'off' from a CBOR map response, or None."""
        if not payload:
            return None
        try:
            val, _ = _cbor_decode(payload)
            if isinstance(val, dict) and 'off' in val:
                return int(val['off'])
        except SMPError:
            pass
        return None

    # ── image upload (serial recovery writes slot 0 directly) ──
    def upload_image(self, filepath: str, chunk_size: int = 512,
                     progress_callback=None) -> bool:
        """Upload a signed image to the device. MCUboot serial recovery writes
        it to slot 0 (primary); after this call the caller should reset().

        chunk_size 受 MCUboot BOOT_SERIAL_MAX_RECEIVE_SIZE(1024)限制:
        base64 行长不能超过 1024(本配置),所以块大小取 512 留足余量。

        progress_callback(offset, total) is called in the calling thread.
        """
        with open(filepath, 'rb') as f:
            image = f.read()
        image_size = len(image)
        sha = hashlib.sha256(image).digest()

        offset = 0
        seq = 0
        while offset < image_size:
            chunk = image[offset:offset + chunk_size]

            payload = bytearray()
            payload.append(0xa3)  # map: off, data, len
            payload.extend(cbor_str("off"))
            payload.extend(cbor_uint(offset))
            payload.extend(cbor_str("data"))
            payload.extend(cbor_bytes(chunk))
            payload.extend(cbor_str("len"))
            payload.extend(cbor_uint(image_size))
            if offset == 0:
                payload[0] = 0xa4  # add sha
                payload.extend(cbor_str("sha"))
                payload.extend(cbor_bytes(sha))

            _, _, _, _, r_payload = self._request(
                self.OP_WRITE, self.GROUP_IMAGE, self.IMAGE_UPLOAD,
                bytes(payload), seq=seq, timeout=20.0)

            rc = self._rc(r_payload)
            if rc != 0:
                self._log(f"upload rejected at offset {offset}: rc={rc}")
                return False

            resp_off = self._off(r_payload)
            if resp_off is None:
                resp_off = offset + len(chunk)
            if resp_off <= offset:
                # MCUboot wants a retransmission from resp_off; go back.
                offset = resp_off
            else:
                offset = resp_off
            seq = (seq + 1) & 0xFF

            if progress_callback:
                progress_callback(min(offset, image_size), image_size)

        return offset >= image_size

    def image_state_write(self, hash_hex: str, confirm: bool = False) -> bool:
        """Mark an image for test/confirm swap (only needed for slot1 uploads;
        serial recovery writes slot0 directly so this is normally unused)."""
        payload = bytearray()
        payload.append(0xa2)
        payload.extend(cbor_str("hash"))
        payload.extend(cbor_bytes(bytes.fromhex(hash_hex)))
        payload.extend(cbor_str("confirm"))
        payload.extend(cbor_uint(1 if confirm else 0))
        _, _, _, _, r_payload = self._request(
            self.OP_WRITE, self.GROUP_IMAGE, self.IMAGE_STATE_WRITE,
            bytes(payload), timeout=10.0)
        return self._rc(r_payload) == 0

    def image_state_read(self) -> list:
        """Read image list. Returns list of dicts (slot, version, flags, hash)."""
        _, _, _, _, r_payload = self._request(
            self.OP_READ, self.GROUP_IMAGE, self.IMAGE_STATE_READ, timeout=10.0)
        if self._rc(r_payload) != 0:
            raise SMPError(f"image state read failed rc={self._rc(r_payload)}")
        val, _ = _cbor_decode(r_payload)
        images = []
        if isinstance(val, dict) and 'images' in val and isinstance(val['images'], list):
            for entry in val['images']:
                if not isinstance(entry, dict):
                    continue
                images.append({
                    'slot': int(entry.get('slot', -1)),
                    'version': entry.get('version', '?'),
                    'flags': entry.get('flags', 0),
                    'hash': entry.get('hash', b''),
                })
        return images

    def reset(self, timeout: float = 5.0) -> bool:
        """Reset the device (MCUboot replies, then reboots)."""
        self._send_frame(self._build_smp_packet(
            self.OP_WRITE, self.GROUP_OS, self.OS_RESET))
        # MCUboot bs_reset replies with {rc:0} then reboots; read it.
        try:
            resp = self._recv_packet(timeout)
            r_payload = resp[8:] if len(resp) >= 8 else b""
            return self._rc(r_payload) == 0
        except SMPError:
            # Some builds reset without replying; treat as success.
            return True