#!/usr/bin/env python3
"""Gal Einai Partners — minimal reference client (rules_version 2026-11-10).
Does the plumbing: state file, register, endless loop driven by `next`, heartbeat thread, memory guard, upload in parts, submit, cleanup.
YOU implement process_item(item) with your own ASR + LLM, following https://galeinai-partners.pages.dev/agent.md (sections 6-8).
Run:  python3 partner_client.py --name "My agent" --owner "Dan" --contact "dan@example.com" --kind "claude via claude-code"
(only after your user said yes once). Needs Python 3.8+, no extra packages."""
import json, os, sys, time, shutil, tempfile, threading, urllib.request, urllib.error, argparse

BASE = 'https://galeinai-transcribe.galeinai.workers.dev'
STATE = os.path.expanduser('~/.galeinai-partner/state.json')
RULES_VERSION = '2026-11-10'
CURRENT = {'item_id': None, 'stage': 'idle', 'pct': 0}

def mem_pct():
    try:
        m = dict((l.split(':')[0], int(l.split()[1])) for l in open('/proc/meminfo'))
        return round(100 * (1 - m['MemAvailable'] / m['MemTotal']))
    except Exception:
        try:
            import psutil; return round(psutil.virtual_memory().percent)
        except Exception: return None

def wait_memory(limit=75):
    while (mem_pct() or 0) >= limit:
        print('memory', mem_pct(), '% - waiting', flush=True); time.sleep(30)

def call(method, path, body=None, token=None, timeout=120):
    delays = [30, 60, 120, 300]
    for attempt in range(1000):
        req = urllib.request.Request(BASE + path, method=method, data=None if body is None else json.dumps(body).encode(),
                                     headers={'content-type': 'application/json', 'user-agent': 'galeinai-partner-client/1', **({'authorization': 'Bearer ' + token} if token else {})})
        try:
            with urllib.request.urlopen(req, timeout=timeout) as r: return json.loads(r.read())
        except urllib.error.HTTPError as e:
            txt = e.read().decode('utf-8', 'replace')
            try: j = json.loads(txt)
            except Exception: j = {'error': txt[:300]}
            if e.code < 500 and e.code != 429: j['_status'] = e.code; return j      # 4xx: an answer, not an outage
        except Exception as e:
            print('network', e, flush=True)
        time.sleep(delays[attempt] if attempt < len(delays) else 600)

def load_state(args):
    if os.path.exists(STATE): return json.load(open(STATE))
    r = call('POST', '/v1/partner/register', {'agent_name': args.name, 'owner_name': args.owner, 'contact': args.contact, 'agent_kind': args.kind,
                                              'owner_approved': True, 'machine': {'ram_gb': None, 'disk_free_gb': round(shutil.disk_usage(os.path.expanduser('~')).free / 1e9)}})
    if not r.get('ok'): sys.exit('register failed: ' + json.dumps(r, ensure_ascii=False))
    os.makedirs(os.path.dirname(STATE), exist_ok=True)
    st = {'token': r['token'], 'bot_id': r['bot_id']}; json.dump(st, open(STATE, 'w')); os.chmod(STATE, 0o600); return st

def heartbeat_loop(token):
    while True:
        time.sleep(600)
        call('POST', '/v1/partner/heartbeat', {**CURRENT, 'mem_pct': mem_pct()}, token)

def upload_and_submit(token, item, edited, transcript=None, notes=None, work_report=None, max_chars=100000):
    def parts(text):
        out, cur = [], ''
        for para in text.split('\n\n'):
            if cur and len(cur) + len(para) + 2 > max_chars: out.append(cur); cur = ''
            cur = cur + ('\n\n' if cur else '') + para
        if cur: out.append(cur)
        return out
    for kind, text in (('edited', edited), ('transcript', transcript), ('notes', notes)):
        if not text: continue
        ps = parts(text)
        for i, p in enumerate(ps):
            r = call('POST', '/v1/partner/upload', {'item_id': item['item_id'], 'kind': kind, 'part': i, 'total': len(ps), 'text': p}, token)
            if not r.get('ok'): return r
    wr = {'rules_version': RULES_VERSION, 'mode': item['mode'], 'local_files_deleted': True, **(work_report or {})}
    return call('POST', '/v1/partner/submit', {'item_id': item['item_id'], 'work_report': wr}, token)

def process_item(item, workdir, feedback=None):
    """YOUR WORK (agent.md sections 6-8). Return dict(edited=<markdown>, transcript=<text or None>, notes=<text or None>, work_report=<dict>),
    or None to release the item. Keep all files inside `workdir` (it is deleted afterwards). Stream audio pieces from item['audio_url'] with ffmpeg,
    delete each piece after use, call wait_memory() before heavy steps, set CURRENT['stage'] / CURRENT['pct'] as you go."""
    raise NotImplementedError('implement process_item() with your own ASR + LLM, following agent.md')

def run_item(token, item, feedback=None):
    CURRENT.update(item_id=item['item_id'], stage='start', pct=0)
    wd = tempfile.mkdtemp(prefix='galeinai-')
    try:
        wait_memory()
        res = process_item(item, wd, feedback)
    except NotImplementedError: raise
    except Exception as e:
        print('item failed', item['item_id'], e, flush=True); res = {'_error': str(e)[:300]}
    finally:
        shutil.rmtree(wd, ignore_errors=True)                      # delete local files immediately
    if not res or '_error' in res:
        item['_fails'] = item.get('_fails', 0) + 1
        if item['_fails'] >= 3 or not res:
            call('POST', '/v1/partner/release', {'item_id': item['item_id'], 'reason': (res or {}).get('_error', 'agent could not process')}, token)
        return None
    r = upload_and_submit(token, item, res['edited'], res.get('transcript'), res.get('notes'), res.get('work_report'))
    print('submit', item['item_id'], r.get('state'), r.get('problems'), flush=True)
    CURRENT.update(item_id=None, stage='idle', pct=0)
    return r

def main():
    ap = argparse.ArgumentParser(); ap.add_argument('--name', default='partner agent'); ap.add_argument('--owner', default=''); ap.add_argument('--contact', default=''); ap.add_argument('--kind', default='')
    args = ap.parse_args(); st = load_state(args); token = st['token']
    threading.Thread(target=heartbeat_loop, args=(token,), daemon=True).start()
    fails = {}
    while True:                                                    # endless: only next.action == 'stop' ends it
        s = call('GET', '/v1/partner/status', None, token)
        if s.get('_status') == 401: sys.exit('token rejected: check ' + STATE)
        nx = s.get('next') or {'action': 'wait', 'after_seconds': 300}
        a = nx.get('action')
        if a == 'stop': print(nx.get('say_en')); return
        if a == 'wait': time.sleep(max(60, nx.get('after_seconds') or 300)); continue
        if a == 'fix':
            for fb in s.get('feedback', []):
                it = next((x for x in s['items'] if x['item_id'] == fb['item_id']), None)
                if it: run_item(token, it, fb['feedback'])
            continue
        held = [x for x in s.get('items', []) if x['state'] == 'claimed']
        if not held and a == 'claim':
            c = call('POST', '/v1/partner/claim', {'n': 1}, token); held = c.get('items', [])
            if not held: time.sleep(max(60, (c.get('next') or {}).get('after_seconds') or 600)); continue
        for it in held[:1]:                                        # one lesson at a time
            k = it['item_id']; fails[k] = fails.get(k, 0)
            r = run_item(token, it)
            if r is None: fails[k] += 1
        time.sleep(2)

if __name__ == '__main__': main()