#!/usr/bin/env python3
"""Pull an image posted into opencode from opencode.db and OCR it via local Unlimited-OCR.

Usage:
  ocr-posted.py [latest | <part_id> | session:<session_id>] [prompt] [save_path]

Defaults: latest image part in DB, prompt "Free OCR."
Env: OPENCODE_DB, OCR_URL
"""
import base64, json, os, sqlite3, sys, urllib.request

DB = os.environ.get(
    'OPENCODE_DB',
    os.path.expanduser('~/.local/share/opencode/opencode.db'))
OCR_URL = os.environ.get('OCR_URL', 'http://127.0.0.1:11436/v1/chat/completions')


def pick(db, arg):
    cur = db.cursor()
    if arg and arg != 'latest':
        if arg.startswith('session:'):
            cur.execute(
                "SELECT id, data FROM part WHERE session_id=? "
                "AND data LIKE '%\"type\":\"file\"%' AND data LIKE '%data:image%' "
                "ORDER BY time_created DESC LIMIT 1",
                (arg[8:],))
        else:
            cur.execute("SELECT id, data FROM part WHERE id=?", (arg,))
    else:
        cur.execute(
            "SELECT id, data FROM part WHERE data LIKE '%\"type\":\"file\"%' "
            "AND data LIKE '%data:image%' ORDER BY time_created DESC LIMIT 1")
    return cur.fetchone()


def ocr(data_url, prompt, max_tokens=4096):
    body = json.dumps({
        "model": "Unlimited-OCR",
        "temperature": 0,
        "repeat_penalty": 1.2,
        "max_tokens": max_tokens,
        "messages": [{"role": "user", "content": [
            {"type": "image_url", "image_url": {"url": data_url}},
            {"type": "text", "text": prompt},
        ]}],
    }).encode()
    req = urllib.request.Request(
        OCR_URL, data=body, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=1800) as r:
        return json.load(r)['choices'][0]['message']['content']


def main():
    arg = sys.argv[1] if len(sys.argv) > 1 else 'latest'
    prompt = sys.argv[2] if len(sys.argv) > 2 else 'Free OCR.'
    save = sys.argv[3] if len(sys.argv) > 3 else None
    db = sqlite3.connect(DB, timeout=10)
    row = pick(db, arg)
    if not row:
        print('no image part found', file=sys.stderr)
        sys.exit(1)
    pid, data = row
    d = json.loads(data)
    url = d.get('url', '')
    if not url.startswith('data:image'):
        print('part %s has no inline image data' % pid, file=sys.stderr)
        sys.exit(1)
    if save:
        open(save, 'wb').write(base64.b64decode(url.split(',', 1)[1]))
        print('saved image to %s' % save, file=sys.stderr)
    print('part: %s  file: %s' % (pid, d.get('filename')), file=sys.stderr)
    print(ocr(url, prompt))


if __name__ == '__main__':
    main()