#!/usr/bin/env python3
"""Compile an Altium **integrated library** (.IntLib) from a .SchLib + .PcbLib, in pure Python.

An .IntLib is an OLE compound file with five streams:
  SchLib/0.schlib      0x02 + zlib(<the .SchLib OLE bytes>)
  PCBLib/0.pcblib      0x02 + zlib(<the .PcbLib OLE bytes>)
  LibCrossRef.Txt      symbol -> footprint cross-reference index (length-prefixed strings)
  Parameters   .bin    the Components-panel search index (parameter string)
  Version.Txt          constant \x00 + uint32(2)

We don't hand-roll a compound-file writer: we take a known-good single-component .IntLib as a
**container template** and swap in this part's five stream contents, then rebuild the OLE with the
same directory tree (reusing schlib_write.read_container / rebuild, already proven to emit
Altium-valid OLEs). The symbol->footprint linkage lives inside the embedded .SchLib itself
(RECORD=45 model link), so the compiled library resolves the footprint with no external PcbLib.
"""
import argparse, json, os, re, struct, sys, zlib
import olefile
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__))))
from schlib_write import read_container

FREESECT = 0xFFFFFFFF; ENDOFCHAIN = 0xFFFFFFFE; FATSECT = 0xFFFFFFFD
SEC = 512; MINI = 64; CUTOFF = 4096


def _pad(b, n):
    return b + b"\x00" * ((-len(b)) % n)


def write_cfb(entries, content):
    """Serialise an OLE/CFB compound file from raw 128-byte directory entries + {sid: bytes},
    reproducing Altium's on-disk SECTOR ORDER: [FAT][directory][miniFAT][mini-stream][big streams].
    (schlib_write.rebuild is spec-valid but places the FAT last; Altium's IntLib reader is strict
    about big streams and expects the conventional FAT-first layout, so we match it here.)
    Streams < 4096 B live in the mini-stream; >= 4096 B get their own FAT sectors."""
    minis = [sid for sid, c in content.items() if len(c) < CUTOFF]
    bigs = [sid for sid, c in content.items() if len(c) >= CUTOFF]

    # mini-stream + mini-FAT
    ministream = b""; mini_start = {}
    for sid in minis:
        mini_start[sid] = len(ministream) // MINI
        ministream += _pad(content[sid], MINI)
    n_mini = len(ministream) // MINI
    minifat = [ENDOFCHAIN] * n_mini
    for sid in minis:
        st = mini_start[sid]; cnt = (len(content[sid]) + MINI - 1) // MINI or 1
        for k in range(cnt - 1):
            minifat[st + k] = st + k + 1
    minifat_bytes = _pad(b"".join(struct.pack("<I", x) for x in minifat), SEC) if n_mini else b""
    n_mf = len(minifat_bytes) // SEC

    ms_p = _pad(ministream, SEC); n_ms = len(ms_p) // SEC
    dir_bytes = _pad(b"".join(bytes(e) for e in entries), SEC); n_dir = len(dir_bytes) // SEC
    big_pad = {sid: _pad(content[sid], SEC) for sid in bigs}
    big_k = {sid: len(big_pad[sid]) // SEC for sid in bigs}

    non_fat = n_dir + n_mf + n_ms + sum(big_k.values())
    n_fat = 1
    while (n_fat + non_fat) * 4 > n_fat * SEC:
        n_fat += 1
    total = n_fat + non_fat
    if n_fat > 109:
        sys.exit("too many FAT sectors for a header-only DIFAT")

    # sector assignment — FAT first, matching Altium
    fat_secs = list(range(0, n_fat)); cur = n_fat
    dir_secs = list(range(cur, cur + n_dir)); cur += n_dir
    mf_secs = list(range(cur, cur + n_mf)); cur += n_mf
    ms_secs = list(range(cur, cur + n_ms)); cur += n_ms
    big_secs = {}
    for sid in bigs:
        big_secs[sid] = list(range(cur, cur + big_k[sid])); cur += big_k[sid]

    # patch directory entries (start sector + 64-bit size)
    def setentry(e, start, size):
        e[116:120] = struct.pack("<I", start); e[120:128] = struct.pack("<Q", size)
    for sid in bigs:
        setentry(entries[sid], big_secs[sid][0], len(content[sid]))
    for sid in minis:
        setentry(entries[sid], mini_start[sid], len(content[sid]))
    setentry(entries[0], ms_secs[0] if n_ms else ENDOFCHAIN, len(ministream))
    dir_bytes = _pad(b"".join(bytes(e) for e in entries), SEC)

    # FAT
    FAT = [FREESECT] * total
    def chain(secs):
        for k in range(len(secs) - 1):
            FAT[secs[k]] = secs[k + 1]
        if secs:
            FAT[secs[-1]] = ENDOFCHAIN
    chain(dir_secs); chain(mf_secs); chain(ms_secs)
    for sid in bigs:
        chain(big_secs[sid])
    for s in fat_secs:
        FAT[s] = FATSECT
    fat_bytes = _pad(b"".join(struct.pack("<I", x) for x in FAT), SEC)

    # header
    h = bytearray(512)
    h[0:8] = bytes([0xD0, 0xCF, 0x11, 0xE0, 0xA1, 0xB1, 0x1A, 0xE1])
    h[24:26] = struct.pack("<H", 0x003E); h[26:28] = struct.pack("<H", 3)
    h[28:30] = struct.pack("<H", 0xFFFE); h[30:32] = struct.pack("<H", 9); h[32:34] = struct.pack("<H", 6)
    h[44:48] = struct.pack("<I", n_fat)
    h[48:52] = struct.pack("<I", dir_secs[0])
    h[56:60] = struct.pack("<I", CUTOFF)
    h[60:64] = struct.pack("<I", mf_secs[0] if n_mf else ENDOFCHAIN)
    h[64:68] = struct.pack("<I", n_mf)
    h[68:72] = struct.pack("<I", ENDOFCHAIN)
    for i in range(109):
        h[76 + 4 * i:80 + 4 * i] = struct.pack("<I", fat_secs[i] if i < n_fat else FREESECT)

    out = bytearray(h) + fat_bytes + dir_bytes + minifat_bytes + ms_p
    for sid in bigs:
        out += big_pad[sid]
    return bytes(out)

SCH_INTERNAL = ":\\SchLib\\0.schlib"
PCB_INTERNAL = ":\\PCBLib\\0.pcblib"
# A known-good single-component .IntLib shipped with the skill, used purely as the OLE container
# skeleton (all five of its streams get overwritten; only its directory tree + the constant
# Version.Txt are reused). Nothing of the template part's data survives into the output.
DEFAULT_TEMPLATE = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                "..", "assets", "templates", "intlib_container.IntLib")


def _astr(s):
    """Altium length-prefixed string: uint32(byte_len+1) + byte(byte_len) + utf8 bytes."""
    b = s.encode("utf-8")
    if len(b) > 254:
        raise ValueError(f"string too long for single-byte length prefix: {s[:40]!r}...")
    return struct.pack("<I", len(b) + 1) + bytes([len(b)]) + b


def build_libcrossref(libref, description, footprint, sch_src, pcb_src, model_type="PCBLIB"):
    out = b"\x00"
    out += struct.pack("<I", 1)
    out += _astr(libref)
    out += _astr(SCH_INTERNAL)
    out += struct.pack("<I", 1)
    out += _astr(description)
    out += _astr(sch_src)
    out += struct.pack("<I", 1)
    out += _astr(footprint)
    out += _astr(model_type)
    out += struct.pack("<I", 1)
    out += _astr(PCB_INTERNAL)
    out += _astr(pcb_src)
    return out


def build_parameters_bin(libref, description, footprint, params, npins, npads):
    fields = ["Comment=*", "Component Kind=Standard", f"Description={description}",
              f"Footprint={footprint}", f"Library Reference={libref}"]
    for k, v in params.items():
        if v not in ("", None):
            fields.append(f"{k}={v}")
    fields += [f"Designator={libref}", "Component Type=Standard", f"Pin Count={npins}    Height=0"]
    sym = "|".join(fields).encode("utf-8")
    out = b"\x00" + struct.pack("<I", len(sym)) + sym
    pad = f"Pad Count={npads}    Height=0".encode("utf-8")
    out += b"\x00" + struct.pack("<I", len(pad) + 1) + pad
    return out


# ---- pull the metadata we need straight out of the embedded .SchLib / .PcbLib ---------------

def sch_facts(schlib_path):
    o = olefile.OleFileIO(schlib_path)
    libref = desc = footprint = None
    params = {}
    npins = 0
    for e in o.listdir(streams=True):
        if e[-1].lower() != "data":
            continue
        t = o.openstream(e).read().decode("utf-8", "replace")
        m = re.search(r"\|RECORD=1\|[^\x00]*?LibReference=([^|]+)", t)
        if not m:
            continue
        libref = m.group(1)
        cd = re.search(r"\|ComponentDescription=([^|]*)", t)
        desc = cd.group(1) if cd else ""
        fp = re.search(r"\|RECORD=45\|[^\x00]*?ModelName=([^|]+)\|ModelType=([^|]+)", t)
        footprint = fp.group(1) if fp else None
        for val, nm in re.findall(r"\|RECORD=41\|[^\x00]*?\|Text=([^|]*)\|Name=([^|]+)\|", t):
            if nm not in ("Comment",):
                params[nm] = val
        npins = len(re.findall(r"\|RECORD=2\|", t))
        break
    o.close()
    return libref, desc or "", footprint, params, npins


def pcb_pad_count(pcblib_path, footprint):
    o = olefile.OleFileIO(pcblib_path)
    n = 0
    try:
        for e in o.listdir(streams=True):
            if e[0].lower() == footprint.lower() and e[-1].lower() == "data":
                b = o.openstream(e).read()
                n = b.upper().count(b"TOP") // 2 or len(re.findall(rb"Pad", b))
    finally:
        o.close()
    return n


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--schlib", required=True)
    ap.add_argument("--pcblib", required=True)
    ap.add_argument("--template", default=DEFAULT_TEMPLATE,
                    help="a known-good single-component .IntLib to reuse as the OLE container "
                         "(defaults to the skill's bundled assets/templates/intlib_container.IntLib)")
    ap.add_argument("--out", required=True)
    ap.add_argument("--pads", type=int, default=0, help="override pad count for the search index")
    a = ap.parse_args()

    # Re-serialise the symbol library through the same FAT-first CFB writer used for the outer
    # container, so the EMBEDDED .schlib has the conventional layout Altium's extractor expects
    # (schlib_write emits a valid but FAT-last OLE; Altium reads that standalone, but its IntLib
    # extractor is stricter — matching the footprint lib's FAT-first layout avoids a read error).
    se, sp, sc = read_container(a.schlib)
    sch = write_cfb(se, sc)
    pcb = open(a.pcblib, "rb").read()
    libref, desc, footprint, params, npins = sch_facts(a.schlib)
    if not libref:
        sys.exit("could not read a LibReference from the .SchLib")
    if not footprint:
        sys.exit("the .SchLib has no RECORD=45 footprint model link — add the footprint model first")
    npads = a.pads or pcb_pad_count(a.pcblib, footprint) or npins

    entries, paths, content = read_container(a.template)
    by = {"/".join(p): sid for sid, p in paths.items()}

    def put(name, data):
        matches = [sid for key, sid in by.items() if key.lower() == name.lower()]
        if not matches:
            sys.exit(f"template IntLib has no stream '{name}'")
        content[matches[0]] = data

    # Altium writes its embedded libs with zlib at DEFAULT level (header 0x789c / FLEVEL=2). Its
    # decompressor rejects other FLEVELs (e.g. level-9's 0x78da) with "Stream read error", so we
    # must match level 6 exactly, not maximise compression.
    put("SchLib/0.schlib", b"\x02" + zlib.compress(sch))
    put("PCBLib/0.pcblib", b"\x02" + zlib.compress(pcb))
    put("LibCrossRef.Txt", build_libcrossref(libref, desc, footprint,
                                             os.path.basename(a.schlib), os.path.basename(a.pcblib)))
    # find the "Parameters   .bin" stream (its name has embedded spaces)
    param_key = next((k for k in by if k.lower().startswith("parameters") and k.lower().endswith(".bin")), None)
    if param_key:
        content[by[param_key]] = build_parameters_bin(libref, desc, footprint, params, npins, npads)

    blob = write_cfb(entries, content)
    open(a.out, "wb").write(blob)

    # ---- validate: re-open, decompress the embedded libs, confirm they are intact -------------
    assert olefile.isOleFile(a.out), "output is not a valid OLE"
    o = olefile.OleFileIO(a.out)
    got = {"/".join(e): o.openstream(e).read() for e in o.listdir(streams=True)}
    o.close()
    sch_back = zlib.decompress(got["SchLib/0.schlib"][1:])
    pcb_back = zlib.decompress(got["PCBLib/0.pcblib"][1:])
    ok = (sch_back == sch and pcb_back == pcb
          and olefile.isOleFile(a.out))
    print(f"wrote {a.out}")
    print(f"  libref={libref!r}  footprint={footprint!r}  desc={desc!r}")
    print(f"  embedded SchLib round-trip: {'OK' if sch_back==sch else 'MISMATCH'} ({len(sch_back)} B)")
    print(f"  embedded PcbLib round-trip: {'OK' if pcb_back==pcb else 'MISMATCH'} ({len(pcb_back)} B)")
    print(f"  LibCrossRef has libref+footprint: "
          f"{libref.encode() in got['LibCrossRef.Txt'] and footprint.encode() in got['LibCrossRef.Txt']}")
    print(f"  streams: {sorted(got)}")
    if not ok:
        sys.exit("VALIDATION FAILED")


if __name__ == "__main__":
    main()