#!/usr/bin/env python3
"""Stamp provenance onto MemPalace drawers: `device` (which machine) and
`agent_kind` (which harness: pi / opencode / miner / cli), plus a *_source key
recording HOW each was determined so an inference is never mistaken for a fact.

Why this exists: mempalace core (3.7.1 through 3.8.0) records neither, and the
single shared bearer token means the server cannot tell clients apart either --
core's `auth_token` is one scalar string (mcp_server.py:7685) and the package
contains zero occurrences of any device/origin concept. The feeder does encode
both facts incidentally -- device in the per-device inbox directory, harness in
the transcript filename prefix (pi_*.jsonl) -- so this recovers them.

Idempotent: only fills keys that are absent, so re-running is cheap and never
downgrades an authoritative value. It must run on a timer, not once: live
re-mining REPLACES metadata rows and silently drops earlier stamps.

RFC 001 7.3.2 places provenance at the sync boundary, never on the agent. As of
2026-08-25 the pi bridge (extensions/pi/mempalace.ts) implements the interim
"edge" row: it defaults added_by/agent/from_agent/created_by to
`<harness>@<device>` from $MEMPALACE_PI_DEVICE, so rule `agent_at_device` below
reads both directly instead of inferring them. This script's remaining job is
(a) historic rows, (b) rows re-mined after a stamp, and (c) diary entries, which
have no writer field at all -- see `diary_host_marker`.

These stamps are self-asserted and therefore ADVISORY (7.3.2 Phase 2): a hint,
never load-bearing for authorization or destructive scoping. When per-device
tokens land (Phase 4) the primary can stamp authoritatively and overwrite with
device_source='token'; the *_source key is what makes that upgrade lossless.

Usage: mempalace-device-stamp [--dry-run]
  --dry-run  report what would be stamped, write nothing. Use this after
             changing any rule below -- the palace is shared by the whole fleet.
"""
import sqlite3, os, re, sys, collections

DRY_RUN = '--dry-run' in sys.argv[1:]
FEED = os.path.expanduser('~/mempalace-feed')
DB = os.path.expanduser('~/.mempalace/palace/chroma.sqlite3')
HOSTPATH = [('/Users/ECSJPER/', 'emb-7kj4vr4g'), ('/Users/joakim/', 'mbp-m1-2020'),
            ('/home/jper/', 'tor-ms22'), ('/home/ecsjper/src', 'synlig')]
DEV_FRAG = [('emb7kj', 'emb-7kj4vr4g'), ('mbp', 'mbp-m1-2020'), ('tor-ms22', 'tor-ms22')]
KNOWN_DEV = set(os.listdir(FEED)) if os.path.isdir(FEED) else set()

base2dev, dupes = {}, set()
if os.path.isdir(FEED):
    for dev in sorted(KNOWN_DEV):
        d = os.path.join(FEED, dev)
        if not os.path.isdir(d): continue
        for f in os.listdir(d):
            if f in base2dev and base2dev[f] != dev: dupes.add(f)
            base2dev[f] = dev

def split_agent(who):
    """'pi@emb-7kj4vr4g' -> ('pi', 'emb-7kj4vr4g'); the forward convention."""
    if who and '@' in who:
        h, d = who.split('@', 1)
        if d in KNOWN_DEV or re.match(r'^[a-z0-9][a-z0-9.-]*$', d or ''):
            return h.strip().lower() or None, d.strip().lower() or None
    return None, None

def resolve_device(sf, added_by, agent, legacy_machine, doc=None):
    for who in (added_by, agent):
        _, d = split_agent(who)
        if d: return d, 'agent_at_device'
    if sf:
        if '/mempalace-feed/' in sf:
            return sf.split('/mempalace-feed/', 1)[1].split('/')[0], 'inbox_path'
        m = re.match(r'^([a-z0-9][a-z0-9-]*):', sf)
        if m and (m.group(1) in KNOWN_DEV): return m.group(1), 'source_prefix'
        b = os.path.basename(sf)
        if b in base2dev and b not in dupes: return base2dev[b], 'basename_index'
        for pref, dev in HOSTPATH:
            if sf.startswith(pref): return dev, 'host_path_heuristic'
    if legacy_machine:
        return legacy_machine.strip().lower(), 'legacy_source_machine'
    if doc:
        # Diary entries have NO writer field: diary_write exposes only
        # agent_name, and a device there would become the wing name
        # (wing = f"wing_{agent_name}"), splitting the diary per host and
        # hiding entries from diary_read. So the pi bridge stamps the device
        # into the entry TEXT as an AAAK field -- `HOST:<device>|SESSION:...`.
        # That is also the only channel a *reader* ever sees, since search
        # projects a fixed key set and diary_read returns content, never
        # metadata. Anchored to a field boundary so it cannot match prose.
        #
        # MUST validate against KNOWN_DEV, for two reasons found by dry-running
        # this rule against the real palace before deploying it:
        #  1. `HOST:` was ALREADY in use in older diary entries with a richer
        #     grammar -- `HOST:emb-7kj4vr4g.f1d3c3f89e3e.v1.8.3.pi0.84.2`
        #     (host.container.image.pi) and bare container ids like
        #     `HOST:2efe2b06f480`. Taking the whole match would invent devices
        #     like "f1d3c3f89e3e.pi0.84.2" and fragment every device query, so
        #     read the first dotted segment and keep it only if it is a real
        #     device. A container id resolves to nothing, which is correct:
        #     containers are not devices and change on every recreate.
        #  2. `HOST:` also carries a DIFFERENT SENSE in some entries -- e.g.
        #     `HOST:exec.via.ssh-controlmaster->alpserv-2(...)` means "the box I
        #     was executing on", not "the box that wrote this". Validation
        #     rejects it, so the two senses cannot be conflated.
        m = re.search(r'(?:^|\|)\s*HOST:\s*([a-z0-9][a-z0-9._-]*)', doc, re.I)
        if m:
            cand = m.group(1).strip().lower().rstrip('.')
            for probe in (cand, cand.split('.')[0]):
                if probe in KNOWN_DEV:
                    return probe, 'diary_host_marker'
    for who in (added_by, agent):
        if who:
            for frag, dev in DEV_FRAG:
                if frag in who.lower(): return dev, 'agent_name_heuristic'
    return None, None

def resolve_kind(sf, added_by, agent, ingest_mode, doc=None):
    for who in (added_by, agent):
        h, _ = split_agent(who)
        if h: return h, 'agent_at_device'
    if sf and '/mempalace-feed/' in sf:
        # Filename shape identifies the harness, but the two shapes differ and a
        # naive "text before the first _" is WRONG for opencode:
        #   pi_<uuid>.jsonl        -> pi
        #   <slug>_ses_<id>.jsonl  -> opencode  (leading segment is an arbitrary slug)
        b = os.path.basename(sf)
        if re.match(r'^pi_', b): return 'pi', 'transcript_prefix'
        if '_ses_' in b: return 'opencode', 'transcript_prefix'
    if doc:
        # Each session's first chunk carries a synthetic header marker.
        m = re.search(r'\|\s*source:\s*(pi|opencode)\b', doc)
        if m: return m.group(1).lower(), 'header_marker'
    for who in (agent, added_by):
        if not who: continue
        w = who.lower()
        if 'opencode' in w: return 'opencode', 'agent_name'
        if w == 'pi' or w.startswith('pi-') or w.startswith('pi_') or 'pi-devbox' in w:
            return 'pi', 'agent_name'
        if w in ('mcp', 'checkpoint'): return 'pi', 'agent_name_weak'
        if w == 'mempalace': return 'miner', 'added_by_default'
    if ingest_mode: return 'miner', 'ingest_mode'
    return None, None

c = sqlite3.connect(DB, timeout=60)
c.execute('pragma busy_timeout=60000')
meta = collections.defaultdict(dict)
for rid, k, sv in c.execute(
        "select id,key,string_value from embedding_metadata where key in "
        "('source_file','added_by','agent','device','agent_kind','ingest_mode','source_machine',"
        "'parent_drawer_id','chroma:document')"):
    meta[rid][k] = sv

dev_rows, kind_rows = [], []
for rid, m in meta.items():
    if not m.get('device'):
        d, how = resolve_device(m.get('source_file'), m.get('added_by'), m.get('agent'),
                               m.get('source_machine'), m.get('chroma:document'))
        if d: dev_rows.append((rid, d, how))
    if not m.get('agent_kind'):
        k, how = resolve_kind(m.get('source_file'), m.get('added_by'), m.get('agent'),
                              m.get('ingest_mode'), m.get('chroma:document'))
        if k: kind_rows.append((rid, k, how))

# --- Propagate within a drawer -------------------------------------------
# Every chunk of one drawer was produced by ONE write call, so they share an
# origin by construction. Without this, any doc-text rule (HOST: marker, header
# marker) resolves only the chunk that happens to contain the marker -- a
# 5-chunk diary entry would end up 1 stamped and 4 blank. Fill-only: never
# overrides a row that resolved on its own evidence, and only propagates when
# the drawer's known chunks agree.
def propagate(rows, key):
    known = {rid: v for rid, v, _ in rows}
    for rid, m in meta.items():
        if m.get(key) and rid not in known:
            known[rid] = m[key]          # already-stamped siblings are evidence too
    by_drawer = collections.defaultdict(set)
    for rid, m in meta.items():
        p = m.get('parent_drawer_id')
        if p and rid in known: by_drawer[p].add(known[rid])
    added = []
    for rid, m in meta.items():
        p = m.get('parent_drawer_id')
        if not p or rid in known or m.get(key): continue
        vals = by_drawer.get(p)
        if vals and len(vals) == 1:      # unanimous, else leave blank
            added.append((rid, next(iter(vals)), 'sibling_chunk'))
    return rows + added

dev_rows = propagate(dev_rows, 'device')
kind_rows = propagate(kind_rows, 'agent_kind')

cur = c.cursor()
def put(rows, key, skey):
    if not rows or DRY_RUN: return
    cur.executemany("insert or replace into embedding_metadata (id,key,string_value) values (?,?,?)",
                    [(r, key, v) for r, v, _ in rows])
    cur.executemany("insert or replace into embedding_metadata (id,key,string_value) values (?,?,?)",
                    [(r, skey, h) for r, _, h in rows])
put(dev_rows, 'device', 'device_source')
put(kind_rows, 'agent_kind', 'agent_kind_source')
if not DRY_RUN: c.commit()
tot = lambda k: c.execute("select count(*) from embedding_metadata where key=?", (k,)).fetchone()[0]
print("%sdevice+%d %s | agent_kind+%d %s | totals device=%d agent_kind=%d" % (
    "DRY RUN would stamp " if DRY_RUN else "stamped ",
    len(dev_rows), dict(collections.Counter(h for _, _, h in dev_rows)) or {},
    len(kind_rows), dict(collections.Counter(v for _, v, _ in kind_rows)) or {},
    tot('device'), tot('agent_kind')))
if DRY_RUN:
    print("   by new rule:", dict(collections.Counter(
        "%s->%s" % (how, v) for _, v, how in dev_rows
        if how in ('diary_host_marker', 'sibling_chunk', 'agent_at_device'))))
