"""Reads a device's app.sab and writes it back out as app.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]
    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 "/"