"""partbin.py -- its bin grows from the whole LDraw catalog (Ivy, Oct 23 2026: "we need to go thru the library and add pieces.
Starting with the most useful ones. And the frog dammit" -> "I should be able to request pieces from the catalog be added to its
bin!!" -> "It should work on its turn, what shape would be right, and what lego does that. We also need descriptions of each piece
somehow").

    catalog(q, n)    search the whole library by meaning: Ivy's words or its wished-for shape
    check(part)      can the engine actually join it? (studs, sockets or holes found in its geometry)
    add(part, who)   check, draw its card pictures and ghost model, describe it, put it in the bin (bin.json, kept forever)
    describe(part)   a plain sentence for a piece: what it looks like and what it's good for (Haiku, from its name + picture)

The base palette lives in play.py; play.PALETTE also takes everything in bin.json. Written by Iris.
"""
import json, os, re, glob, subprocess, threading, time
import numpy as np
HERE = os.path.dirname(os.path.abspath(__file__))
LIB = os.path.expanduser('~/lego/ldraw/parts')
BIN = os.path.join(HERE, 'bin.json')               # [{"part", "who", "when"}]
DESC = os.path.join(HERE, 'parts', 'desc.json')    # {part: sentence}
INDEX = os.path.join(HERE, 'catalog.json')         # [[part, name, keywords]]
_LOCK = threading.RLock()
ID_OK = re.compile(r'^\d+[a-z]?\d?$|^\d+c\d+$')    # 3040b, 33320, 122c01 -- not printed variants (3001p01) or subparts
SKIP = ('~', '=', '_', '|')

def extras():
    try: return [b['part'] for b in json.load(open(BIN))]
    except (OSError, ValueError): return []

def index():
    """every usable part in the library: id, name, keywords (cached in catalog.json)"""
    if os.path.exists(INDEX): return json.load(open(INDEX))
    out = []
    for f in sorted(glob.glob(os.path.join(LIB, '*.dat'))):
        pid = os.path.basename(f)[:-4]
        if not ID_OK.match(pid): continue
        try: head = [l.strip() for _, l in zip(range(25), open(f, errors='ignore'))]
        except OSError: continue
        name = head[0][2:].strip() if head and head[0].startswith('0 ') else ''
        if not name or name.startswith(SKIP) or 'Sticker' in name or 'Moved to' in name or 'Obsolete' in name or 'Minifig' in name: continue
        kw = ' '.join(l.split('!KEYWORDS', 1)[1].strip() for l in head if '!KEYWORDS' in l)
        cat = ' '.join(l.split('!CATEGORY', 1)[1].strip() for l in head if '!CATEGORY' in l)
        out.append([pid, re.sub(r'\s+', ' ', name), (cat + ' ' + kw).strip()])
    json.dump(out, open(INDEX, 'w')); return out

_E = None
def _emb():
    """MiniLM vectors for the whole catalog (cached as catalog_emb.npy)"""
    global _E
    if _E is not None: return _E
    import sys; sys.path.insert(0, os.path.join(HERE, 'teacher')); import teach
    ix = index(); path = os.path.join(HERE, 'catalog_emb.npy')
    if os.path.exists(path) and len(np.load(path)) == len(ix): _E = np.load(path)
    else:
        _E = np.asarray(teach.emb([f"{n}. {k}" for _, n, k in ix]), dtype=np.float32); np.save(path, _E)
    _E = _E / np.maximum(np.linalg.norm(_E, axis=1, keepdims=True), 1e-9)
    return _E

_WARMING = []
def _warm():
    """build the catalog's meaning vectors in a background process on spare CPU (it took minutes in-process under load)"""
    if _WARMING and _WARMING[0].poll() is None: return
    _WARMING[:] = [subprocess.Popen(['flock', '-n', '/var/tmp/brick_loop/catalog_emb.lock', 'chrt', '--idle', '0', 'python3', '-c',
                                     'import partbin; partbin._emb()'], cwd=HERE, env={**os.environ, 'HIP_VISIBLE_DEVICES': '', 'CUDA_VISIBLE_DEVICES': ''},
                                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)]
def catalog(q, n=12, offset=0):
    """the parts that best match q: every word of q found in the name first, then by meaning (page by offset)"""
    want = n; n = n + offset
    ix = index(); q = (q or '').strip()
    if not q: return []
    words = [w for w in re.findall(r'[a-z0-9]+', q.lower())]
    exact = [i for i, (pid, nm, kw) in enumerate(ix) if words and all(w in (pid + ' ' + nm + ' ' + kw).lower() for w in words)]
    try: rank = json.load(open(os.path.join(HERE, 'common_rank.json')))     # how many sets use it (Brick Architect, all years)
    except (OSError, ValueError): rank = {}
    in_name = lambda i: all(w in (ix[i][0] + ' ' + ix[i][1]).lower() for w in words)
    exact.sort(key=lambda i: (not in_name(i), rank.get(ix[i][0], 10**6), len(ix[i][1])))   # name matches, then the most-used parts
    out = exact[:n]
    if len(out) < n and not os.path.exists(os.path.join(HERE, 'catalog_emb.npy')):
        _warm()                                          # meaning-search comes online once its vectors are built (minutes, once)
    elif len(out) < n:
        import sys; sys.path.insert(0, os.path.join(HERE, 'teacher')); import teach
        v = np.asarray(teach.emb([q])[0], dtype=np.float32); v /= max(np.linalg.norm(v), 1e-9)
        for i in np.argsort(-(_emb() @ v)):
            if len(out) >= n: break
            if int(i) not in out: out.append(int(i))
    have = set(extras())
    import play
    return [{'part': ix[i][0], 'name': ix[i][1], 'in_bin': ix[i][0] in have or ix[i][0] in play.PALETTE} for i in out[offset:offset + want]]

PREV = os.path.join(HERE, 'parts', 'prev')
_PREV_LOCK = threading.Semaphore(2)                  # at most two LeoCAD renders at once
def preview(part):
    """a picture of any catalog part (its card picture if it's in the bin; otherwise drawn once and kept). Path or None."""
    part = re.sub(r'[^0-9a-z]', '', str(part).lower())
    for p in (os.path.join(HERE, 'parts', f'{part}.png'), os.path.join(PREV, f'{part}.png')):
        if os.path.exists(p): return p
    if not os.path.exists(os.path.join(LIB, part + '.dat')): return None
    os.makedirs(PREV, exist_ok=True); ldr = os.path.join(PREV, f'{part}.ldr')
    with _PREV_LOCK:
        if not os.path.exists(ldr[:-4] + '.png'):
            open(ldr, 'w').write(f'0 {part}\n1 14 0 0 0 1 0 0 0 1 0 0 0 1 {part}.dat\n')
            subprocess.run(['flatpak', 'run', '--filesystem=home', 'org.leocad.LeoCAD', '-l', os.path.expanduser('~/lego/ldraw'), '-i',
                            ldr[:-4] + '.png', '-w', '160', '-h', '160', '--viewpoint', 'home', ldr], capture_output=True, timeout=120)
    return ldr[:-4] + '.png' if os.path.exists(ldr[:-4] + '.png') else None

def check(part):
    """None if the engine can join it, else why not"""
    import play
    if not os.path.exists(os.path.join(LIB, part + '.dat')): return 'not in the library'
    try: p = play.G.fresh(part)
    except Exception as e: return f'could not load it ({str(e)[:60]})'
    if not (len(p.studs) or len(p.sockets) or len(p.holes) or play.new_snaps(part)):     # bars, clips, axles, hinges count too
        return 'no studs, sockets, holes, bars, clips, axles or hinges on it, so nothing could join it'
    if max(p.mesh.extents) > 160: return 'too big for the bin'
    return None

def _render(part):
    os.makedirs(os.path.join(HERE, 'parts'), exist_ok=True)
    for tag, col in (('', 14), ('_w', 15)):
        ldr = os.path.join(HERE, 'parts', f'{part}{tag}.ldr')
        open(ldr, 'w').write(f'0 {part}\n1 {col} 0 0 0 1 0 0 0 1 0 0 0 1 {part}.dat\n')
        subprocess.run(['flatpak', 'run', '--filesystem=home', 'org.leocad.LeoCAD', '-l', os.path.expanduser('~/lego/ldraw'), '-i',
                        ldr[:-4] + '.png', '-w', '160', '-h', '160', '--viewpoint', 'home', ldr], capture_output=True, timeout=120)
    import to_glb; to_glb.part_glb(part, os.path.join(HERE, 'parts', f'{part}.glb'))

def _desc_all():
    try: return json.load(open(DESC))
    except (OSError, ValueError): return {}

def describe(part, name=None, force=False):
    """what it looks like and what it's good for, in plain words a kid could follow (Haiku, from the name and the white render)"""
    d = _desc_all()
    if part in d and not force: return d[part]
    import imagine as I                               # Gemini first, Haiku second (Oct 24 2026: seeding 30 parts, Haiku 429'd every one)
    name = name or next((nm for pid, nm, _ in index() if pid == part), part)
    img = os.path.join(HERE, 'parts', f'{part}_w.png')
    t = ''
    for wait in (0, 20, 60):                          # the subscription rate-limits (429) when other sessions are busy
        time.sleep(wait)
        try:
            t, _ = I.small_ask(f'This is the Lego part "{name}" (LDraw {part}), pictured. In ONE plain sentence of at most 22 words, say what it '
                         f'looks like (shape, size in studs, anything special like a slope, hole, curve or face) and what builders use it for. '
                         f'No part number. Answer with the sentence only.', [img] if os.path.exists(img) else [], 120, prefer='gemini')
            break
        except Exception as e: print('describe', part, 'failed:', str(e)[:60], flush=True)
    t = (t or '').strip().strip('"')
    if t:
        with _LOCK:
            d = _desc_all(); d[part] = t; json.dump(d, open(DESC, 'w'), indent=0)
    return t

def add(part, who='Ivy', describe_it=True):
    """put a catalog part in its bin for good. Returns (ok, message)"""
    import play
    part = str(part).strip().lower()
    if part in play.PALETTE: return False, 'already in its bin'
    why = check(part)
    if why: return False, why
    with _LOCK:
        _render(part)
        name = next((nm for pid, nm, _ in index() if pid == part), part)
        try: names = json.load(open(os.path.join(HERE, 'parts', 'names.json')))
        except (OSError, ValueError): names = {}
        names[part] = name; json.dump(names, open(os.path.join(HERE, 'parts', 'names.json'), 'w'))
        b = json.load(open(BIN)) if os.path.exists(BIN) else []
        if part not in [x['part'] for x in b]: b.append({'part': part, 'who': who, 'when': time.strftime('%Y-%m-%d %H:%M')})
        json.dump(b, open(BIN, 'w'), indent=1)
        play.PALETTE.append(part); play.CX.NAMES[part] = name
    try: describe_it and describe(part, name)         # never let a description block a piece; describe_missing() fills gaps
    except Exception as e: print('describe failed:', e, flush=True)
    return True, f'{name} is in its bin now'

def describe_missing(pause=2.0):
    """a description for every piece in its bin that doesn't have one yet (slowly: one call every couple of seconds)"""
    import play
    d = _desc_all(); done = 0
    for part in play.PALETTE:
        if part in d: continue
        if not os.path.exists(os.path.join(HERE, 'parts', f'{part}_w.png')): continue
        if describe(part, play.CX.NAMES.get(part)): done += 1
        time.sleep(pause)
    return done

def ldraw_id(pid):
    """Brick Architect / BrickLink numbers -> the LDraw file (3069 -> 3069b)"""
    for c in (pid, pid + 'b', pid + 'a', pid + 'c'):
        if os.path.exists(os.path.join(LIB, c + '.dat')): return c
    return None