"""Append XR Agent samples to NDJSON. Python 3 standard library only.
Usage: python xr_agent_collect.py SESSION_ID [output.ndjson]
Reconnects with replay and deduplicates persisted IDs. Ctrl+C stops.
"""
import json
import sys
import time
import urllib.request
import urllib.error
from pathlib import Path

BASE = 'https://rag.ingevision.cloud/xr-agent/public/api'
if len(sys.argv) < 2:
    raise SystemExit('Usage: python xr_agent_collect.py SESSION_ID [output.ndjson]')
session = sys.argv[1]
if len(session) != 32 or any(c not in '0123456789abcdef' for c in session):
    raise SystemExit('Invalid session ID; retrieve IDs from ' + BASE + '/sessions')
target = Path(sys.argv[2] if len(sys.argv) > 2 else session + '.ndjson')
seen = set()
if target.exists():
    with target.open(encoding='utf-8') as existing:
        for line in existing:
            record = json.loads(line)  # Fail visibly if an existing file is damaged.
            if record.get('id'):
                seen.add(record['id'])
try:
    with target.open('a', encoding='utf-8') as out:
        while True:
            try:
                with urllib.request.urlopen(BASE + '/sessions/' + session + '/stream.ndjson', timeout=40) as response:
                    for raw in response:
                        record = json.loads(raw)
                        if record.get('type') == 'heartbeat':
                            continue
                        key = record.get('id')
                        if key and key not in seen:
                            out.write(json.dumps(record, ensure_ascii=False) + '\n')
                            out.flush()
                            seen.add(key)
            except urllib.error.HTTPError as error:
                if error.code in (400, 401, 403, 404):
                    raise SystemExit(str(error))
                print(error, file=sys.stderr)
            except (urllib.error.URLError, TimeoutError, ConnectionError) as error:
                print(error, file=sys.stderr)
            time.sleep(3)
except KeyboardInterrupt:
    print('Stopped. Saved to ' + str(target))