#!/usr/bin/env python3
# Python KEKS implementation
# Copyright (C) 2024-2025 Sergey Matveev <stargrave@stargrave.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, version 3 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""Python KEKS encoder/decoder implementation

`KEKS <http://www.keks.cypherpunks.su>`__ is compact, deterministic,
concise and streaming binary serialisation format. It is aimed to be
lightweight in terms of CPU, memory, storage and codec implementation
size usage. It supports wide range of data types, making it able to
transparently replace JSON.

It has :py:func:`loads` and :py:func:`dumps` functions, similar to
native :py:module:`json` library's. KEKS supports dictionaries, lists,
None, booleans, UUID, IPv6 addresses, floats (currently not fully
implemented!), integers (including big ones), datetime, UTF-8 and
binary strings.

There is special :py:func:`keks.Raw` namedtuple, that holds arbitrary
KEKS encoded data, that can not be represented in native Python types.
Also there is :py:func:`keks.Blob` namedtuple, that holds the data, that
is encoded streamingly in chunks.
"""

from collections import namedtuple
from datetime import datetime
from datetime import timedelta
from datetime import timezone
from ipaddress import IPv6Address
from math import ceil as _ceil
from math import isnan as _isnan
from uuid import UUID


TagEOC = 0x00
TagNIL = 0x01
TagFalse = 0x02
TagTrue = 0x03
TagHexlet = 0x04
TagList = 0x08
TagMap = 0x09
TagBlob = 0x0B
TagPInt = 0x0C
TagNInt = 0x0D
TagFloatNaN = 0x10
TagFloatPinf = 0x11
TagFloatNinf = 0x12
TagFloat = 0x13
TagTAI64 = 0x18
TagTAI64N = 0x19
TagTAI64NA = 0x1A
TagMagic = 0x4B
TagStr = 0x80
TagUTF8 = 0x40
NaN = float("nan")


def _byte(v):
    return v.to_bytes(1, "big")


TagEOCb = _byte(TagEOC)
TagNILb = _byte(TagNIL)
TagFalseb = _byte(TagFalse)
TagTrueb = _byte(TagTrue)
TagHexletb = _byte(TagHexlet)
TagListb = _byte(TagList)
TagMapb = _byte(TagMap)
TagBlobb = _byte(TagBlob)
TagPIntb = _byte(TagPInt)
TagNIntb = _byte(TagNInt)
TagTAI64b = _byte(TagTAI64)
TagTAI64Nb = _byte(TagTAI64N)
TagMagicb = _byte(TagMagic)
TagFloatNaNb = _byte(TagFloatNaN)
TagFloatPinfb = _byte(TagFloatPinf)
TagFloatNinfb = _byte(TagFloatNinf)
TagFloatb = _byte(TagFloat)


class DecodeError(ValueError):
    pass


class NotEnoughData(DecodeError):
    def __init__(self, n):
        self.n = n
        super().__init__()

    def __str__(self):
        return "%d bytes expected" % self.n

    def __repr__(self):
        return "%s(%s)" % (self.__class__.__name__, self)


class Raw:
    __slots__ = ("v",)

    def __init__(self, v: bytes):
        self.v = v

    def __bytes__(self) -> bytes:
        return self.v

    def __eq__(self, other) -> bool:
        if not isinstance(other, self.__class__):
            return False
        return self.v == other.v

    def __repr__(self) -> str:
        return "Raw(%s)" % self.v.hex()


class Magic:
    __slots__ = ("v",)

    def __init__(self, v: bytes):
        if len(v) > 12:
            raise ValueError("too long")
        self.v = v

    def __bytes__(self) -> bytes:
        return self.v

    def __eq__(self, other) -> bool:
        if not isinstance(other, self.__class__):
            return False
        return self.v == other.v

    def __repr__(self) -> str:
        return "Magic(%r)" % self.v


class Hexlet:
    __slots__ = ("v",)

    def __init__(self, v):
        if isinstance(v, UUID):
            self.v = v.bytes
        elif isinstance(v, IPv6Address):
            self.v = v.packed
        else:
            if len(v) != 16:
                raise ValueError("wrong hexlet len")
            self.v = v

    def __bytes__(self) -> bytes:
        return self.v

    def __eq__(self, other) -> bool:
        if not isinstance(other, self.__class__):
            return False
        return self.v == other.v

    def __repr__(self) -> str:
        return "HEXLET(%s)" % self.asUUID()

    def asUUID(self) -> UUID:
        return UUID(bytes=self.v)

    def asIP(self) -> IPv6Address:
        return IPv6Address(self.v)


class Float:
    __slots__ = ("m", "e")

    def __init__(self, m: int, e: int):
        self.m = m
        self.e = e

    def __eq__(self, other) -> bool:
        if not isinstance(other, self.__class__):
            return False
        return (self.m == other.m) and (self.e == other.e)

    def __float__(self) -> float:
        return self.m * pow(2, self.e)

    def __repr__(self) -> str:
        return "Float(m=%d, e=%d)" % (self.m, self.e)


Blob = namedtuple("Blob", ("l", "v"))


def LenFirstUTF8Sort(v):
    v = v.encode("utf-8")
    return (len(v), v)


TAI64Base = 0x4000000000000000
Leapsecs1972 = 10
Leapsecs = tuple(
    Leapsecs1972 + i + int(
        datetime(y, m, 1, 0, 0, 0, 0, tzinfo=timezone.utc).timestamp()
    )
    for i, (y, m) in enumerate((
        (1972, 7),
        (1973, 1),
        (1974, 1),
        (1975, 1),
        (1976, 1),
        (1977, 1),
        (1978, 1),
        (1979, 1),
        (1980, 1),
        (1981, 7),
        (1982, 7),
        (1983, 7),
        (1985, 7),
        (1988, 1),
        (1990, 1),
        (1991, 1),
        (1992, 7),
        (1993, 7),
        (1994, 7),
        (1996, 1),
        (1997, 7),
        (1999, 1),
        (2006, 1),
        (2009, 1),
        (2012, 7),
        (2015, 7),
        (2017, 1),
    ))
)


def _str(v, utf8):
    l = len(v)
    if l >= (63 + ((1 << 8)-1) + ((1 << 16)-1)):
        lv = 63
        l -= (lv + ((1 << 8)-1) + ((1 << 16)-1))
        lb = l.to_bytes(8, "big")
    elif l >= (62 + ((1 << 8)-1)):
        lv = 62
        l -= (lv + ((1 << 8)-1))
        lb = l.to_bytes(2, "big")
    elif l >= 61:
        lv = 61
        l -= lv
        lb = l.to_bytes(1, "big")
    else:
        lv = l
        lb = b""
    t = TagStr
    if utf8:
        t |= TagUTF8
    return _byte(t | lv) + lb + v


def utc2tai(secs):
    """Add leapseconds to UTC
    """
    secs += Leapsecs1972
    for leapsec in Leapsecs:
        if secs >= leapsec:
            secs += 1
    return secs


def dumps(v):
    if v is None:
        return TagNILb
    if v is False:
        return TagFalseb
    if v is True:
        return TagTrueb
    if isinstance(v, UUID) or isinstance(v, IPv6Address):
        return dumps(Hexlet(v))
    if isinstance(v, Hexlet):
        return TagHexletb + v.v
    if isinstance(v, float):
        if v == 0.0:
            return dumps(Float(0, 0))
        if _isnan(v):
            return TagFloatNaNb
        if v == float("+inf"):
            return TagFloatPinfb
        if v == float("-inf"):
            return TagFloatNinfb
        raise NotImplementedError("no full FLOAT support")
    if isinstance(v, Float):
        return TagFloatb + dumps(v.m) + dumps(v.e)
    if isinstance(v, datetime):
        ms = v.microsecond
        v = v.replace(microsecond=0)
        secs = int(v.replace(tzinfo=timezone.utc).timestamp())
        secs = utc2tai(secs) + TAI64Base
        if ms == 0:
            return TagTAI64b + secs.to_bytes(8, "big")
        return (
            TagTAI64Nb + secs.to_bytes(8, "big") +
            (ms * 1000).to_bytes(4, "big")
        )
    if isinstance(v, Raw):
        return v.v
    if isinstance(v, bytes):
        return _str(v, utf8=False)
    if isinstance(v, str):
        return _str(v.encode("utf-8"), utf8=True)
    if isinstance(v, int):
        t = TagPIntb
        if v < 0:
            t = TagNIntb
            v = (-v) - 1
        if v == 0:
            return t + dumps(b"")
        return t + dumps(v.to_bytes(_ceil(v.bit_length() / 8), "big"))
    if isinstance(v, Blob):
        assert (v.l > 0) and (v.l <= (1 << 64))
        l, v = v.l, v.v
        raws = [TagBlobb, (l-1).to_bytes(8, "big")]
        append = raws.append
        chunks = len(v) // l
        for i in range(chunks):
            append(dumps(v[i*l:(i+1)*l]))
        left = len(v) - chunks*l
        assert left < l
        append(dumps(b"") if (left == 0) else dumps(v[-left:]))
        return b"".join(raws)
    if isinstance(v, (list, tuple)):
        return b"".join([TagListb] + [dumps(i) for i in v] + [TagEOCb])
    if isinstance(v, set):
        if not all(isinstance(i, str) for i in v):
            raise ValueError("set can contain only strings")
        return dumps({i: None for i in v})
    if isinstance(v, dict):
        raws = [TagMapb]
        append = raws.append
        keys = v.keys()
        if not all(isinstance(k, str) for k in keys):
            raise ValueError("map keys can be only strings")
        keys = sorted(keys, key=LenFirstUTF8Sort)
        if (len(keys) > 0) and len(keys[0]) == 0:
            raise ValueError("map keys can not be empty")
        for k in keys:
            append(dumps(k))
            append(dumps(v[k]))
        append(TagEOCb)
        return b"".join(raws)
    if isinstance(v, Magic):
        return b"".join((b"KEKS", v.v, b"\x00" * (12 - len(v.v))))
    raise NotImplementedError("unsupported type")


def _int(v):
    s, tail = _loads(v, _allowContainers=False)
    if not isinstance(s, bytes):
        raise DecodeError("non-BIN in INT")
    if s == b"":
        return 0, tail
    if s[0] == 0:
        raise DecodeError("non-minimal encoding")
    return int.from_bytes(s, "big"), tail


_EOC = object()
_tais = {TagTAI64: 8, TagTAI64N: 12, TagTAI64NA: 16}


def tai2utc(secs, leapsecUTCAllow=False):
    """Subtract leapseconds from TAI

    :returns: UTC seconds, or None of TAI equals to leap second
    """
    diff = Leapsecs1972
    for leapsec in Leapsecs:
        if secs < leapsec:
            break
        diff += 1
        if secs == leapsec:
            if leapsecUTCAllow:
                break
            return None
    return secs - diff


def _loads(v, sets=False, leapsecUTCAllow=False, _allowContainers=True):
    if len(v) == 0:
        raise NotEnoughData(1)
    b = v[0]
    if b == TagEOC:
        return _EOC, v[1:]
    if b == TagNIL:
        return None, v[1:]
    if b == TagFalse:
        return False, v[1:]
    if b == TagTrue:
        return True, v[1:]
    if b == TagHexlet:
        if len(v) < 1+16:
            raise NotEnoughData(1+16-len(v))
        return Hexlet(v[1:1+16]), v[1+16:]
    if b == TagFloatNaN:
        return NaN, v[1:]
    if b == TagFloatPinf:
        return float("+inf"), v[1:]
    if b == TagFloatNinf:
        return float("-inf"), v[1:]
    if b == TagFloat:
        m, v = _loads(v[1:])
        if type(m) != int:
            raise DecodeError("non-int m float")
        e, v = _loads(v)
        if type(e) != int:
            raise DecodeError("non-int e float")
        if (m == 0) and (e == 0):
            return 0.0, v
        return Float(m, e), v
    l = _tais.get(b)
    if l is not None:
        if len(v) < 1+l:
            raise NotEnoughData(1+l - len(v))
        secs = int.from_bytes(v[1:1+8], "big")
        if secs >= (1 << 63):
            raise DecodeError("reserved TAI64 value is in use")
        nsecs = 0
        if l > 8:
            nsecs = int.from_bytes(v[1+8:1+8+4], "big")
            if (l == 12) and (nsecs == 0):
                raise DecodeError("non-minimal TAI64N")
            if nsecs > 999999999:
                raise DecodeError("too many nanoseconds")
        asecs = 0
        if l > 12:
            asecs = int.from_bytes(v[1+8+4:1+8+4+4], "big")
            if asecs == 0:
                raise DecodeError("non-minimal TAI64NA")
            if asecs > 999999999:
                raise DecodeError("too many attoseconds")
        secs = tai2utc(secs - TAI64Base, leapsecUTCAllow)
        if secs is None:
            return Raw(v[:1+l]), v[1+l:]
        if (asecs > 0) or ((nsecs % 1000) > 0):
            # Python can represent neither attoseconds, nor nanoseconds
            return Raw(v[:1+l]), v[1+l:]
        try:
            dt = datetime(1970, 1, 1) + timedelta(seconds=secs)
        except OverflowError:
            return Raw(v[:1+l]), v[1+l:]
        dt += timedelta(microseconds=nsecs // 1000)
        return dt, v[1+l:]
    if (b & TagStr) > 0:
        l = b & 0b00111111
        if l < 61:
            llen = 0
        elif l == 61:
            llen = 1
        elif l == 62:
            llen = 2
            l += ((1 << 8)-1)
        elif l == 63:
            llen = 8
            l += ((1 << 8)-1) + ((1 << 16)-1)
        if llen > 0:
            if len(v) < 1+llen:
                raise NotEnoughData(1+llen-len(v))
            l += int.from_bytes(v[1:1+llen], "big")
        if len(v) < 1+llen+l:
            raise NotEnoughData(1+llen+l-len(v))
        s = v[1+llen:1+llen+l]
        if (b & TagUTF8) > 0:
            try:
                s = s.decode("utf-8")
            except UnicodeDecodeError as err:
                raise DecodeError("invalid UTF-8") from err
            if "\x00" in s:
                raise DecodeError("null byte in UTF-8")
        return s, v[1+llen+l:]
    if b == TagPInt:
        return _int(v[1:])
    if b == TagNInt:
        i, v = _int(v[1:])
        return (-1 - i), v
    if not _allowContainers:
        raise DecodeError("unknown tag")
    if b == TagList:
        ret = []
        v = v[1:]
        while True:
            i, v = _loads(v, sets=sets, leapsecUTCAllow=leapsecUTCAllow)
            if i == _EOC:
                break
            ret.append(i)
        return ret, v
    if b == TagMap:
        ret = {}
        v = v[1:]
        kPrev = b""
        allNILs = True
        while True:
            k, v = _loads(v, _allowContainers=False)
            if k == _EOC:
                break
            if not isinstance(k, str):
                raise DecodeError("non-string key")
            kUTF8 = k.encode("utf-8")
            if (
                (len(kUTF8) < len(kPrev)) or
                ((len(kUTF8) == len(kPrev)) and (kUTF8 <= kPrev))
            ):
                if len(k) == 0:
                    raise DecodeError("empty key")
                raise DecodeError("unsorted keys")
            i, v = _loads(v, sets=sets, leapsecUTCAllow=leapsecUTCAllow)
            if i == _EOC:
                raise DecodeError("unexpected EOC")
            ret[k] = i
            kPrev = kUTF8
            if i is not None:
                allNILs = False
        if sets and allNILs:
            ret = set(ret.keys())
        return ret, v
    if b == TagBlob:
        if len(v) < 1+8:
            raise NotEnoughData(1+8-len(v))
        l = 1 + int.from_bytes(v[1:1+8], "big")
        v = v[1+8:]
        raws = []
        while True:
            i, v = _loads(v, _allowContainers=False)
            if not isinstance(i, bytes):
                raise DecodeError("unexpected tag")
            if len(i) == l:
                raws.append(i)
            elif len(i) < l:
                raws.append(i)
                break
            else:
                raise DecodeError("wrong chunk len")
        return Blob(l, b"".join(raws)), v
    if b == TagMagic:
        if len(v) < 16:
            raise NotEnoughData(16-len(v))
        if v[:4] != b"KEKS":
            raise DecodeError("wrong magic")
        return Magic(v[4:16].rstrip(b"\x00")), v[16:]
    raise DecodeError("unknown tag")


def loads(v, **kwargs):
    """Decode KEKS-encoded data.

    :param bool sets: transform maps with NIL-only values to set()s
    :param bool leapsecUTCAllow: allow TAI64 values equal to leap seconds,
                                 to be decoded as datetime UTC value. Raw()
                                 value will be returned instead
    :returns: decoded data and the undecoded tail
    :rtype: (any, bytes)
    """
    try:
        ret, tail = _loads(v, **kwargs)
    except RecursionError as err:
        raise DecodeError("deep recursion") from err
    if ret == _EOC:
        raise DecodeError("unexpected EOC")
    return ret, tail


if __name__ == "__main__":
    from argparse import ArgumentParser
    parser = ArgumentParser(description="Decode KEKS file")
    parser.add_argument(
        "--nosets", action="store_true",
        help="Do not fold NIL-valued maps to sets")
    parser.add_argument(
        "--leapsec-utc-allow", action="store_true",
        help="Allow leapseconded UTC values")
    args = parser.parse_args()
    from sys import stdin
    data = stdin.buffer.read()
    obj = Magic(b"")
    from pprint import pprint
    while isinstance(obj, Magic):
        obj, data = loads(
            data,
            sets=not args.nosets,
            leapsecUTCAllow=args.leapsec_utc_allow is True,
        )
        pprint(obj)
    if data != b"":
        print("tail:", data.hex())
