#!/usr/bin/env python3
"""Write Vecmocon's mandatory parameters directly INTO an Altium schematic-library (.SchLib)
file, in pure Python — no Altium needed.

Reading a value out of a .SchLib is a plain stream read (see altium_refs.py). Writing is harder:
parameters live as text records inside the component's ``Data`` stream, and adding/removing them
changes that stream's size, which means rebuilding the whole OLE compound container around it.
This script does exactly that while **preserving every original byte except the Data stream(s)
it must change** — all other streams (FileHeader, Storage, …), the directory tree, the Altium
CLSIDs and timestamps are copied through untouched, and only the size/offset fields that must
move are recomputed.

What it does to the target component's Data stream:
  - removes the Ultra-Librarian default params ``Manufacturer_Name`` / ``Manufacturer_Part_Number``
    (they duplicate the SOP ``Manufacturer`` / ``Manufacturer Part`` — override with --keep or a
    "remove" list in params.json),
  - adds/updates the SOP parameters from params.json (see references/schlib_parameters.md),
  - leaves pins, graphics, the Comment and all other records exactly as they were.

Usage:
  python schlib_write.py --schlib IN.SchLib --params params.json --out OUT.SchLib

params.json (same shape altium_params.py uses):
  {"component":"JMK105BJ105KV-F",           # LibRef / component-storage name; omit -> all comps
   "parameters":{"Value":"1u","Manufacturer":"Taiyo Yuden", ...},
   "remove":["Manufacturer_Name","Manufacturer_Part_Number"]}   # optional; this is the default

IMPORTANT: this writes Altium's own binary format from outside Altium. It is validated to
re-open as a well-formed OLE with every other stream byte-identical, but ALWAYS open the result
in Altium once to confirm it loads before relying on it.
"""
import argparse, json, os, struct, sys, hashlib
import olefile

FREESECT=0xFFFFFFFF; ENDOFCHAIN=0xFFFFFFFE; FATSECT=0xFFFFFFFD
SEC=512; MINI=64; CUTOFF=4096
DEFAULT_REMOVE=["Manufacturer_Name", "Manufacturer_Part_Number"]


def le16(b,o): return struct.unpack('<H',b[o:o+2])[0]
def le32(b,o): return struct.unpack('<I',b[o:o+4])[0]
def pad(b,n): return b+b'\x00'*((-len(b))%n)


# ----------------------------------------------------------------- typeid template columns

# Template columns that are internal library bookkeeping, NOT symbol parameters.
NON_PARAM_COLS = {"MPN_make_type", "Skill Version", "Template Version",
                  "Library Ref", "Library Path", "Footprint Ref", "Footprint Path"}


def template_param_names(template_path, typeid):
    """The symbol-parameter columns for a typeid = every column on that typeid's template sheet
    EXCEPT the internal bookkeeping ones (the tag, the two version columns, and the four design
    Ref/Path columns). These are the engineering parameters that belong on the symbol."""
    import openpyxl
    wb = openpyxl.load_workbook(template_path, read_only=True)
    if typeid not in wb.sheetnames:
        raise SystemExit(f"no template sheet for typeid '{typeid}'")
    ws = wb[typeid]
    return [ws.cell(1, c).value for c in range(1, ws.max_column + 1)
            if ws.cell(1, c).value and ws.cell(1, c).value not in NON_PARAM_COLS]


# ----------------------------------------------------------------- read the container

def read_container(path):
    """Return (raw_dir_entries[list of bytearray(128)], sid->path, sid->content-bytes)."""
    raw=open(path,'rb').read()
    dir_start=le32(raw,48)
    difat=[le32(raw,76+4*i) for i in range(109)]
    fat=[]
    for s in difat:
        if s>=0xFFFFFFFC: continue
        sec=raw[512+s*SEC:512+(s+1)*SEC]
        fat+=[struct.unpack('<I',sec[k:k+4])[0] for k in range(0,SEC,4)]
    def chain(start):
        out=[]; s=start
        while s not in (ENDOFCHAIN,FREESECT) and s<len(fat):
            out.append(s); s=fat[s]
        return out
    dirbytes=b''.join(raw[512+s*SEC:512+(s+1)*SEC] for s in chain(dir_start))
    entries=[bytearray(dirbytes[i:i+128]) for i in range(0,len(dirbytes),128)]

    def name_of(e):
        nlen=le16(e,64)
        return e[0:max(0,nlen-2)].decode('utf-16-le') if nlen>=2 else ''
    paths={}
    def walk(sid, prefix):
        if sid>=0xFFFFFFFC or sid>=len(entries): return
        e=entries[sid]
        walk(le32(e,68), prefix)                       # left
        nm=name_of(e); paths[sid]=prefix+[nm]
        if e[66]==1:                                   # storage -> recurse children
            walk(le32(e,76), prefix+[nm])
        walk(le32(e,72), prefix)                       # right
    walk(le32(entries[0],76), [])                      # from root's child

    ole=olefile.OleFileIO(path)
    bypath={'/'.join(e):ole.openstream(e).read() for e in ole.listdir(streams=True)}
    ole.close()
    content={sid:bypath['/'.join(p)] for sid,p in paths.items() if '/'.join(p) in bypath}
    return entries, paths, content


# ----------------------------------------------------------------- edit a Data stream

def _uid(name):
    h=''.join(c for c in hashlib.md5(name.encode()).hexdigest().upper() if c.isalpha())
    return (h+"ABCDEFGH")[:8]


def _param_record(idx, name, value):
    s=(f"|RECORD=41|IndexInSheet={idx}|OwnerPartId=1|Justification=4|FontID=2|IsHidden=T"
       f"|Text={value}|Name={name}|UniqueID={_uid(name)}")
    payload=s.encode('utf-8')+b'\x00'                 # utf-8: handles Ω, µ, °, ± in names/values
    return struct.pack('<I',len(payload))+payload


def _leading_text_records(data):
    """Split Data into (list of leading text-record byte-blocks, tail bytes). The leading run is
    the contiguous |RECORD=…| text records at the start (RECORD=1 + the parameter block); it
    stops at the first non-text record (the binary pin records), which — with everything after —
    is returned untouched as the tail."""
    recs=[]; i=0
    while i+4<=len(data):
        L=struct.unpack('<I',data[i:i+4])[0]
        if L==0 or i+4+L>len(data): break
        payload=data[i+4:i+4+L]
        if payload[-1:]!=b'\x00' or not payload[:8].startswith(b'|RECORD='):
            break
        recs.append(data[i:i+4+L]); i+=4+L
    return recs, data[i:]


def _rec_name(block):
    t=block[4:-1].decode('utf-8','replace')
    return t.split('|Name=')[1].split('|')[0] if '|Name=' in t else None


def _patch_field(block, field, value):
    """Replace |field=...| inside a length-prefixed text record, re-framing its 4-byte length.
    Used to set the component's ComponentDescription in the RECORD=1 header."""
    import re
    text = block[4:-1].decode('utf-8', 'replace')
    if f"|{field}=" in text:
        text = re.sub(rf"\|{re.escape(field)}=[^|]*", f"|{field}={value}", text, count=1)
    elif text.startswith("|RECORD="):
        text = text + f"|{field}={value}"
    payload = text.encode('utf-8') + b'\x00'
    return struct.pack('<I', len(payload)) + payload


def edit_data(data, params, remove):
    """Return a new Data stream: drop `remove` params, drop any SOP-name params (re-added
    fresh), keep everything else, then append the SOP params. Also mirror the `Description`
    parameter into the component's ComponentDescription field (the Altium 'Description' shown in
    the component properties). Pins/graphics/tail untouched."""
    leading, tail = _leading_text_records(data)
    sop_names=set(params)
    desc = params.get("Description")
    kept=[]
    for blk in leading:
        nm=_rec_name(blk)
        if nm is not None and (nm in remove or nm in sop_names):
            continue                                   # drop UL duplicates + stale SOP copies
        if desc is not None and blk[4:-1].startswith(b"|RECORD=1|"):
            blk = _patch_field(blk, "ComponentDescription", desc)   # component Description field
        kept.append(blk)
    added=[_param_record(20+i, nm, val) for i,(nm,val) in enumerate(params.items())]
    return b''.join(kept)+b''.join(added)+tail


# ----------------------------------------------------------------- rebuild the container

def _set(e, off, val, n): e[off:off+n]=val.to_bytes(n,'little')

def rebuild(entries, content):
    """Rebuild an OLE compound file from raw directory entries + {sid: content}. Small streams
    (<4096) go in the mini-stream; big streams get their own FAT sectors. Directory-entry bytes
    are preserved; only each entry's start-sector and size are patched."""
    minis=[sid for sid,c in content.items() if len(c)<CUTOFF]
    bigs =[sid for sid,c in content.items() if len(c)>=CUTOFF]

    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

    # sector layout: [big streams][ministream][miniFAT][directory][FAT]
    ms_p=pad(ministream,SEC); n_ms=len(ms_p)//SEC
    big_secs={}; cur=0
    for sid in bigs:
        b=pad(content[sid],SEC); k=len(b)//SEC
        big_secs[sid]=(list(range(cur,cur+k)), b); cur+=k
    ms_secs=list(range(cur,cur+n_ms)); cur+=n_ms
    mf_secs=list(range(cur,cur+n_mf)); cur+=n_mf
    dir_bytes=pad(b''.join(entries),SEC); n_dir=len(dir_bytes)//SEC
    dir_secs=list(range(cur,cur+n_dir)); cur+=n_dir
    data_secs=cur
    n_fat=1
    while (data_secs+n_fat)*4 > n_fat*SEC: n_fat+=1
    fat_secs=list(range(cur,cur+n_fat)); total=cur+n_fat
    if n_fat>109: sys.exit("too many FAT sectors for a header-only DIFAT (file unexpectedly large)")

    # patch directory entries: start-sector + size
    for sid in bigs:  _set(entries[sid],116,big_secs[sid][0][0],4); entries[sid][120:128]=struct.pack('<Q',len(content[sid]))
    for sid in minis: _set(entries[sid],116,mini_start[sid],4);     entries[sid][120:128]=struct.pack('<Q',len(content[sid]))
    entries[0][120:128]=struct.pack('<Q',len(ministream))            # root: ministream size
    _set(entries[0],116, ms_secs[0] if n_ms else ENDOFCHAIN, 4)
    dir_bytes=pad(b''.join(entries),SEC)

    # FAT
    FAT=[FREESECT]*total
    def ch(secs):
        for k in range(len(secs)-1): FAT[secs[k]]=secs[k+1]
        if secs: FAT[secs[-1]]=ENDOFCHAIN
    for sid in bigs: ch(big_secs[sid][0])
    ch(ms_secs); ch(mf_secs); ch(dir_secs)
    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)
    for sid in bigs: out+=big_secs[sid][1]
    out+=ms_p+minifat_bytes+dir_bytes+fat_bytes
    return bytes(out)


# ----------------------------------------------------------------- driver

def write_params(schlib, params_json, out, typeid=None, template=None):
    component=params_json.get("component") or None
    fields=dict(params_json.get("parameters", {}) or {})
    remove=params_json.get("remove", DEFAULT_REMOVE)

    # If a typeid+template are given, guarantee the FULL template parameter set is written:
    # every engineering column for that typeid becomes a symbol parameter (value from the
    # params if provided, else blank). This is what makes every .SchLib carry the complete,
    # consistent parameter set the template defines — not just whatever was hand-listed.
    typeid = typeid or params_json.get("typeid")
    template = template or params_json.get("template")
    if typeid and template:
        for name in template_param_names(template, typeid):
            fields.setdefault(name, "")
    entries, paths, content = read_container(schlib)

    # target Data stream sid(s): a stream named 'Data' whose parent storage == component (or all)
    targets=[]
    for sid,p in paths.items():
        if len(p)>=2 and p[-1].lower()=="data" and sid in content:
            if component is None or p[-2]==component:
                targets.append(sid)
    if not targets:
        sys.exit(f"no component Data stream found" + (f" for '{component}'" if component else ""))
    for sid in targets:
        content[sid]=edit_data(content[sid], fields, remove)

    blob=rebuild(entries, content)
    open(out,'wb').write(blob)

    # self-check: re-open and confirm it's a valid OLE with the params present
    if not olefile.isOleFile(out):
        sys.exit("ERROR: rebuilt file is not a valid OLE — aborting")
    ole=olefile.OleFileIO(out)
    ok=True
    for e in ole.listdir(streams=True):
        if e[-1].lower()=="data":
            t=ole.openstream(e).read().decode('utf-8','ignore')
            for nm in fields:
                if f"|Name={nm}|" not in t and f"|Name={nm}\x00" not in t and f"Name={nm}" not in t:
                    ok=False
    ole.close()
    n_comp=len(targets)
    print(f"wrote {out}  ({len(fields)} params into {n_comp} component(s); removed {remove}) "
          f"{'[self-check OK]' if ok else '[WARN: verify params]'}")


def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--schlib", required=True)
    ap.add_argument("--params", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--typeid", help="component typeid; with --template, writes that typeid's "
                                     "full template parameter set (blank where not provided)")
    ap.add_argument("--template", help="path to template.xlsx (defaults to the skill's)")
    a=ap.parse_args()
    template = a.template
    if a.typeid and not template:
        template = os.path.join(os.path.dirname(__file__), "..", "assets", "template", "template.xlsx")
    write_params(a.schlib, json.load(open(a.params, encoding="utf-8")), a.out,
                 typeid=a.typeid, template=template)


if __name__ == "__main__":
    main()