"""Reads a device's app.sab and writes it back out as app.sax, and encodes an app.sax
into app.sab (encode_sax).

Ported from Py_sab2sax (sax.py/sab.py), itself a port of sedona.offline.OfflineApp
(decodeAppBinary + encodeAppXml); its output matches `sedonac app.sab` byte for byte,
down to the `<!-- /path -->` comments and Java's float formatting. Kit manifests come
from Sedona home, found the way the editor finds them everywhere else.

A .sab holds config props only, so runtime values cannot come back: like sedonac, the
manifest default stands in for them, and a value equal to its default is not written.
"""

import base64
import math
import struct
import xml.etree.ElementTree as ET

import schema


class SabError(Exception):
    pass


class _Reader:
    def __init__(self, data):
        self.b, self.p = data, 0

    def take(self, n):
        if self.p + n > len(self.b):
            raise SabError(f"truncated sab at offset {self.p}")
        v = self.b[self.p:self.p + n]
        self.p += n
        return v

    def u1(self): return self.take(1)[0]
    def u2(self): return struct.unpack(">H", self.take(2))[0]
    def i4(self): return struct.unpack(">i", self.take(4))[0]
    def i8(self): return struct.unpack(">q", self.take(8))[0]
    def f4(self): return struct.unpack(">f", self.take(4))[0]
    def f8(self): return struct.unpack(">d", self.take(8))[0]

    def str_(self):
        end = self.b.index(0, self.p)
        s = self.b[self.p:end].decode("ascii")
        self.p = end + 1
        return s


# --- schema ---------------------------------------------------------------------------

class Slot:
    __slots__ = ("name", "type", "flags", "default", "facets")

    def __init__(self, elem):
        self.name = elem.get("name")
        self.type = elem.get("type")
        self.flags = elem.get("flags", "")
        self.default = elem.get("default")
        self.facets = {f.get("name"): f.get("val") for f in elem.iter()
                       if f is not elem and f.get("name") is not None}


def _load_kits(parts):
    """[(kitName, checksumHex)] in kit-id order -> (types by kit, types by qname)."""
    kits, types_by_q = [], {}
    for name, cks in parts:
        path = schema.find_manifest_path(name, cks)
        if path is None:
            raise SabError(f"no manifest for kit '{name}' in Sedona home")
        mroot = ET.parse(path).getroot()
        tlist = sorted(mroot.findall("type"), key=lambda t: int(t.get("id")))
        for t in tlist:
            types_by_q[f"{name}::{t.get('name')}"] = t
        kits.append(tlist)
    return kits, types_by_q


def _resolve_slots(qname, types_by_q, cache):
    """The type's slots, base chain first, indexed by slot id."""
    if qname in cache:
        return cache[qname]
    if qname not in types_by_q:
        raise SabError(f"type {qname} is not in the app's kits -- a kit build that "
                       f"needs kit '{qname.split('::')[0]}', which the app does not list?")
    t = types_by_q[qname]
    base = t.get("base")
    out = list(_resolve_slots(base, types_by_q, cache)) if base else []
    by_name = {s.name: i for i, s in enumerate(out)}
    for elem in sorted(t.findall("slot"), key=lambda s: int(s.get("id"))):
        slot = Slot(elem)
        if slot.name in by_name:           # an action override keeps the inherited id
            out[by_name[slot.name]] = slot
        else:
            by_name[slot.name] = len(out)
            out.append(slot)
    cache[qname] = out
    return out


# --- values ---------------------------------------------------------------------------

def _int(raw):   # Integer.decode: supports 0x
    if raw in (None, "", "null"):
        return 0
    return int(raw, 0) if raw.lower().startswith(("0x", "-0x")) else int(raw)


def _flt(raw):
    if raw == "null":
        return float("nan")             # explicit null only
    if raw in (None, ""):
        return 0.0                      # Value.defaultForType -> zero
    return float(raw)


def _decode_value(r, vtype, flags):
    if "s" in flags:                                   # asStr Buf: u2 len (incl NUL)
        n = r.u2()
        return r.take(n).split(b"\0")[0].decode("ascii", errors="replace")
    if vtype == "bool":
        v = r.u1()
        return None if v == 2 else (v == 1)
    if vtype == "byte":   return r.u1()
    if vtype == "short":  return r.u2()
    if vtype == "int":    return r.i4()
    if vtype == "long":   return r.i8()
    if vtype == "float":  return r.f4()
    if vtype == "double": return r.f8()
    if vtype in ("sys::Buf", "buf"):
        return bytes(r.take(r.u2()))
    raise SabError(f"unhandled slot type {vtype}")


def _default_value(slot):
    """The manifest default, or Value.defaultForType's zero."""
    vtype, flags, raw = slot.type, slot.flags, slot.default
    if "s" in flags:
        return raw or ""
    if vtype == "bool":
        return None if raw == "null" else raw in ("true", "1")
    if vtype in ("byte", "short", "int", "long"):
        return _int(raw)
    if vtype == "float":
        # the default is a 32-bit float too: 0.1 as a double is not the file's 0.1
        return _to_f32(_flt(raw))
    if vtype == "double":
        return _flt(raw)
    if vtype in ("sys::Buf", "buf"):
        return base64.b64decode(raw) if raw else b""
    raise SabError(f"unhandled slot type {vtype}")


def _same(a, b):
    if isinstance(a, float) and isinstance(b, float):
        return (math.isnan(a) and math.isnan(b)) or a == b
    return type(a) is type(b) and a == b


# --- Java number formatting -----------------------------------------------------------

_F32_MAX = 3.4028234663852886e38
_F32_INF_EDGE = 3.4028235677973366e38     # midpoint to 2**128: below rounds to F32_MAX


def _to_f32(v):
    try:
        return struct.unpack(">f", struct.pack(">f", v))[0]
    except OverflowError:
        if abs(v) < _F32_INF_EDGE:
            return math.copysign(_F32_MAX, v)
        return math.copysign(math.inf, v)


def _shortest(a, single):
    """Fewest digits that still round-trip -- what Java's toString picks."""
    for prec in range(0, 17):
        s = "%.*e" % (prec, a)
        v = float(s)
        if single:
            v = _to_f32(v)
        if v == a:
            break
    mant, exp = s.split("e")
    digits = mant.replace(".", "").rstrip("0") or "0"
    return digits, int(exp)


def _java_float(v, single=True):
    """Float/Double.toString: always a decimal point, E-form outside [1e-3, 1e7)."""
    if math.isnan(v):
        return "NaN"
    if math.isinf(v):
        return "-Infinity" if v < 0 else "Infinity"
    neg = math.copysign(1.0, v) < 0
    a = abs(v)
    if a == 0.0:
        return "-0.0" if neg else "0.0"
    digits, exp = _shortest(a, single)
    if -3 <= exp < 7:
        if exp >= 0:
            head = digits[:exp + 1].ljust(exp + 1, "0")
            tail = digits[exp + 1:] or "0"
        else:
            head, tail = "0", "0" * (-exp - 1) + digits
        out = head + "." + tail
    else:
        out = digits[0] + "." + (digits[1:] or "0") + "E" + str(exp)
    return "-" + out if neg else out


def _encode_string(val, slot):
    """Value.encodeString()."""
    vtype, flags = slot.type, slot.flags
    if "s" in flags:
        return _literal(val)
    if vtype == "bool":
        return "null" if val is None else ("true" if val else "false")
    if vtype == "float":
        return "null" if math.isnan(val) else _java_float(val, True)
    if vtype == "double":
        return "null" if math.isnan(val) else _java_float(val, False)
    if vtype in ("sys::Buf", "buf"):
        return base64.b64encode(val).decode("ascii")
    return str(val)


def _literal(s):
    """TextUtil.toLiteral."""
    return "".join({"\0": "\\0", "\n": "\\n", "\r": "\\r", "\t": "\\t",
                    "\\": "\\\\", '"': '\\"', "$": "\\$"}.get(c, c) for c in s)


def _safe(s):
    """XWriter.safe: escape markup, quotes, and anything outside printable ASCII."""
    out = []
    for c in s:
        o = ord(c)
        if o < 0x20 or o > 0x7e or c in "'\"":
            out.append("&#x%x;" % o)
        elif c == "<":
            out.append("&lt;")
        elif c == ">":
            out.append("&gt;")
        elif c == "&":
            out.append("&amp;")
        else:
            out.append(c)
    return "".join(out)


# --- the app --------------------------------------------------------------------------

class SabApp:
    """A decoded app.sab: kits, components by id, and links.

    Each component is a dict: id, type, name, parent, kids, and props -- a list of
    (Slot, value) for every config prop, defaults included.
    """

    def __init__(self, data):
        r = _Reader(data)
        if r.i4() != 0x73617070:
            raise SabError("not a .sab (bad magic)")
        ver = r.i4()
        if ver != 0x0003:
            raise SabError(f"unsupported sab version 0x{ver:x}")

        self.kit_parts = []
        for _ in range(r.u1()):
            name = r.str_()
            self.kit_parts.append((name, f"{r.i4() & 0xFFFFFFFF:08x}"))
        kits, self._types_by_q = _load_kits(self.kit_parts)
        self._cache = {}
        r.u2()                                  # maxId, recomputed on the way out

        qname_of = {}
        for ki, ((kname, _), tlist) in enumerate(zip(self.kit_parts, kits)):
            for t in tlist:
                qname_of[(ki, int(t.get("id")))] = f"{kname}::{t.get('name')}"

        by_id = {}
        while True:
            cid = r.u2()
            if cid == 0xFFFF:
                break
            ki, ti = r.u1(), r.u1()
            if (ki, ti) not in qname_of:
                raise SabError(f"comp {cid}: no type {ti} in kit {ki}")
            qname = qname_of[(ki, ti)]
            c = {"id": cid, "type": qname, "name": r.str_(), "kids": [], "props": []}
            c["parentId"], c["childId"], c["siblingId"] = r.u2(), r.u2(), r.u2()
            for slot in self.slots(qname):
                if "a" in slot.flags or "c" not in slot.flags:
                    continue                    # config props only, in slot-id order
                c["props"].append((slot, _decode_value(r, slot.type, slot.flags)))
            if r.u1() != ord(";"):
                raise SabError(f"corrupted component {cid} {c['name']}")
            if cid in by_id:
                raise SabError(f"duplicate id: {cid}")
            by_id[cid] = c

        def must(cid, what, owner):
            if cid not in by_id:
                raise SabError(f"comp {owner['id']} {owner['name']}: missing {what} {cid}")
            return by_id[cid]

        for c in by_id.values():
            c["parent"] = must(c["parentId"], "parent", c) if c["parentId"] != 0xFFFF else None
            if c["childId"] != 0xFFFF:
                kid = must(c["childId"], "child", c)
                while True:
                    c["kids"].append(kid)
                    if kid["siblingId"] == 0xFFFF:
                        break
                    kid = must(kid["siblingId"], "nextSibling", kid)

        self.links = []
        while True:
            fc = r.u2()
            if fc == 0xFFFF:
                break
            fs, tc, ts = r.u1(), r.u2(), r.u1()
            for cid in (fc, tc):
                if cid not in by_id:
                    raise SabError(f"link references unknown comp id {cid}")
            self.links.append((by_id[fc], fs, by_id[tc], ts))
        if r.u1() != ord("."):
            raise SabError("invalid app end marker")

        if 0 not in by_id:
            raise SabError("no app component (id 0)")
        self.comps = by_id
        self.app = by_id[0]

    def slots(self, qname):
        return _resolve_slots(qname, self._types_by_q, self._cache)

    def kit_names(self):
        return [name for name, _ in self.kit_parts]

    def prop(self, comp, name, default=None):
        for slot, val in comp["props"]:
            if slot.name == name:
                return val
        return default

    def slot(self, comp, name):
        return next((s for s in self.slots(comp["type"]) if s.name == name), None)

    # --- writing app.sax ---

    def to_sax(self):
        out = []
        w = out.append
        w("<?xml version='1.0'?>\n")
        w("<sedonaApp>\n")
        w("<schema>\n")
        for name, cks in self.kit_parts:
            w(f'  <kit name="{_safe(name)}" checksum="{cks}" />\n')
        w("</schema>\n")
        w("<app>\n")
        self._write_props(w, self.app, 2)
        for kid in self.app["kids"]:
            self._write_comp(w, kid, 2)
        w("</app>\n")
        w("<links>\n")
        for fc, fs, tc, ts in self.links:
            w('  <link from="%s.%s" to="%s.%s"/>\n' % (
                _safe(_path_of(fc)), _safe(self._slot_name(fc, fs)),
                _safe(_path_of(tc)), _safe(self._slot_name(tc, ts))))
        w("</links>\n")
        w("</sedonaApp>\n")
        return "".join(out)

    def _slot_name(self, comp, slot_id):
        slots = self.slots(comp["type"])
        if slot_id >= len(slots):
            raise SabError(f"{comp['type']} has no slot id {slot_id}")
        return slots[slot_id].name

    def _write_props(self, w, c, indent):
        """Only props whose value differs from the slot default get written."""
        n = 0
        for slot, val in c["props"]:
            if _same(val, _default_value(slot)):
                continue
            w('%s<prop name="%s" val="%s"/>\n' % (
                " " * indent, _safe(slot.name), _safe(_encode_string(val, slot))))
            n += 1
        return n

    def _write_comp(self, w, c, indent):
        pad = " " * indent
        w(f"{pad}<!-- {_path_of(c)} -->\n")
        w(f'{pad}<comp name="{_safe(c["name"])}" id="{c["id"]}" type="{_safe(c["type"])}"')
        body = []
        n = self._write_props(body.append, c, indent + 2)
        for kid in c["kids"]:
            self._write_comp(body.append, kid, indent + 2)
        if not n and not c["kids"]:
            w("/>\n")
        else:
            w(">\n")
            w("".join(body))
            w(f"{pad}</comp>\n")


def _path_of(c):
    parts = []
    while c["parent"] is not None:
        parts.append(c["name"])
        c = c["parent"]
    return "/" + "/".join(reversed(parts)) if parts else "/"


# --- encoding app.sax to app.sab ------------------------------------------------------
#
# Ported from Py_sab2sax's sab.py (sedona.offline.OfflineApp decodeAppXml +
# encodeAppBinary), with the value parsing brought in line with the Java Value classes:
# asStr text is unescaped as TextUtil.fromLiteral does, "nullbool"/"nullfloat" are
# read, and integers take Integer.decode's '#' hex and leading-zero octal.

def _decode_int(raw):
    """Integer.decode."""
    if raw in (None, "", "null"):
        return 0
    s = raw.strip()
    neg = s.startswith("-")
    if s[:1] in "+-":
        s = s[1:]
    if s[:2].lower() == "0x":
        v = int(s[2:], 16)
    elif s.startswith("#"):
        v = int(s[1:], 16)
    elif len(s) > 1 and s.startswith("0"):
        v = int(s[1:], 8)
    else:
        v = int(s)
    return -v if neg else v


def _from_literal(s):
    """TextUtil.fromLiteral."""
    if "\\" not in s:
        return s
    escapes = {"0": "\0", "n": "\n", "r": "\r", "t": "\t", "\\": "\\", '"': '"', "$": "$"}
    out, i = [], 0
    while i < len(s):
        c = s[i]
        if c == "\\":
            i += 1
            if i >= len(s) or s[i] not in escapes:
                raise SabError(f"unknown escape sequence in {s!r}")
            c = escapes[s[i]]
        out.append(c)
        i += 1
    return "".join(out)


class _Writer:
    def __init__(self):
        self.b = bytearray()

    def u1(self, v): self.b.append(v & 0xFF)
    def u2(self, v): self.b += struct.pack(">H", v & 0xFFFF)
    def i4(self, v): self.b += struct.pack(">i", ((v + 2**31) % 2**32) - 2**31)
    def i8(self, v): self.b += struct.pack(">q", v)
    def f4(self, v): self.b += struct.pack(">f", _to_f32(v))
    def f8(self, v): self.b += struct.pack(">d", v)

    def str_(self, s):
        self.b += s.encode("ascii")
        self.u1(0)


def _encode_value(out, slot, raw):
    """One config prop, from its sax text (or the manifest default when absent)."""
    vtype, flags = slot.type, slot.flags
    if "s" in flags:                                   # asStr Buf: u2 len (incl NUL)
        data = _from_literal(raw or "").encode("ascii")
        out.u2(len(data) + 1)
        out.b += data
        out.u1(0)
    elif vtype == "bool":
        if raw in ("null", "nullbool"):
            out.u1(2)
        elif raw in (None, "", "false", "true"):
            out.u1(1 if raw == "true" else 0)
        else:
            raise SabError(f"invalid bool {raw!r}")
    elif vtype == "byte":
        out.u1(_decode_int(raw))
    elif vtype == "short":
        out.u2(_decode_int(raw))
    elif vtype == "int":
        out.i4(_decode_int(raw))
    elif vtype == "long":
        out.i8(_decode_int(raw))
    elif vtype in ("float", "double"):
        v = float("nan") if raw in ("null", "nullfloat", "nulldouble") else _flt(raw)
        (out.f4 if vtype == "float" else out.f8)(v)
    elif vtype in ("sys::Buf", "buf"):
        data = base64.b64decode(raw) if raw else b""
        out.u2(len(data))
        out.b += data
    else:
        raise SabError(f"unhandled slot type {vtype}")


def encode_sax(sax, checksums=None, kits=None):
    """app.sax (text or bytes) -> app.sab bytes, as `sedonac app.sax` writes it.

    checksums, {kit name: checksum hex}, stands in for the sax's own -- how an app goes
    to a device whose kits are other builds than the ones it was saved with. kits,
    [(name, checksum hex)], replaces the schema altogether, as Ontrol's PutApp rewrites
    it for the kits chosen to put. Every kit's manifest must be in Sedona home under
    the checksum used.
    """
    root = ET.fromstring(sax)
    schema_elem, app_elem = root.find("schema"), root.find("app")
    if schema_elem is None or app_elem is None:
        raise SabError("not a Sedona app: no <schema> or <app>")
    if kits is None:
        kits = [(k.get("name"), (checksums or {}).get(k.get("name")) or k.get("checksum"))
                for k in schema_elem.findall("kit")]
    parts = []
    for name, cks in kits:
        if not cks:
            raise SabError(f"no checksum for kit '{name}'")
        parts.append((name, f"{int(cks, 16):08x}"))
    parts.sort(key=lambda p: ("" if p[0] == "sys" else "\x01" + p[0]))   # sortKits
    for name, cks in parts:
        path = schema.find_manifest_path(name, cks)
        if path is None or not path.endswith(f"-{cks}.xml"):
            raise SabError(f"no manifest for kit '{name}' {cks} in Sedona home")
    kits, types_by_q = _load_kits(parts)
    type_ids = {}
    for ki, ((kname, _), tlist) in enumerate(zip(parts, kits)):
        for t in tlist:
            type_ids[f"{kname}::{t.get('name')}"] = (ki, int(t.get("id")))
    cache = {}

    comps = []

    def walk(elem, qtype, parent):
        if qtype not in type_ids:
            raise SabError(f"unknown type '{qtype}' ({elem.get('name')})")
        c = {"elem": elem, "type": qtype, "parent": parent, "kids": [],
             "id": int(elem.get("id")) if elem.get("id") else None}
        comps.append(c)
        if parent is not None:
            parent["kids"].append(c)
        for kid in elem.findall("comp"):
            walk(kid, kid.get("type"), c)
        return c

    app = walk(app_elem, "sys::App", None)
    app["id"] = 0
    used = {c["id"] for c in comps if c["id"] is not None}
    next_id = 1
    for c in comps:                                 # ids for comps the sax gives none
        if c["id"] is None:
            while next_id in used:
                next_id += 1
            c["id"] = next_id
            used.add(next_id)
    by_id = {}
    for c in comps:
        if c["id"] in by_id:
            raise SabError(f"duplicate id: {c['id']}")
        by_id[c["id"]] = c

    out = _Writer()
    out.i4(0x73617070)                              # "sapp"
    out.i4(0x0003)                                  # version 0.3
    out.u1(len(parts))
    for name, cks in parts:
        out.str_(name)
        out.i4(int(cks, 16))
    out.u2(max(by_id))

    for cid in sorted(by_id):
        c = by_id[cid]
        ki, ti = type_ids[c["type"]]
        out.u2(cid)
        out.u1(ki)
        out.u1(ti)
        out.str_(c["elem"].get("name", "app") if cid else "app")
        out.u2(c["parent"]["id"] if c["parent"] else 0xFFFF)
        out.u2(c["kids"][0]["id"] if c["kids"] else 0xFFFF)
        sibling = 0xFFFF
        if c["parent"]:
            kids = c["parent"]["kids"]
            i = kids.index(c)
            if i + 1 < len(kids):
                sibling = kids[i + 1]["id"]
        out.u2(sibling)
        props = {p.get("name"): p.get("val") for p in c["elem"].findall("prop")}
        for slot in _resolve_slots(c["type"], types_by_q, cache):
            if "a" in slot.flags or "c" not in slot.flags:
                continue                            # config props only
            _encode_value(out, slot, props.get(slot.name, slot.default))
        out.u1(ord(";"))
    out.u2(0xFFFF)

    def ref(text):
        path, slot_name = text.rsplit(".", 1)
        cur = app
        for part in (p for p in path.split("/") if p):
            cur = next((k for k in cur["kids"] if k["elem"].get("name") == part), None)
            if cur is None:
                raise SabError(f"link to a missing component: {text}")
        names = [s.name for s in _resolve_slots(cur["type"], types_by_q, cache)]
        if slot_name not in names:
            raise SabError(f"link to a missing slot: {text}")
        return cur["id"], names.index(slot_name)

    links = root.find("links")
    for link in (links.findall("link") if links is not None else ()):
        fc, fs = ref(link.get("from"))
        tc, ts = ref(link.get("to"))
        out.u2(fc)
        out.u1(fs)
        out.u2(tc)
        out.u1(ts)
    out.u2(0xFFFF)
    out.u1(ord("."))
    return bytes(out.b)