#!/usr/bin/env python3
"""Small SQLite exact-key memory workload; no embeddings, vendor API, or secrets."""
import concurrent.futures
import json
import math
import os
import platform
import sqlite3
import tempfile
import time
from pathlib import Path

NOW = 2000000000  # Fixed fixture clock, not wall time.


def retrieve(conn, tenant, item_id):
    if tenant not in {'tenant-0', 'tenant-1', 'tenant-2', 'tenant-3'}:
        raise ValueError('Authenticated tenant binding required')
    return conn.execute(
        'SELECT body, source FROM memory WHERE tenant=? AND id=? AND expires>? AND trusted=1',
        (tenant, item_id, NOW)).fetchone()


def seed(path, count):
    conn = sqlite3.connect(path)
    conn.execute('PRAGMA journal_mode=WAL')
    conn.execute('CREATE TABLE memory(tenant TEXT, id INTEGER, body TEXT, source TEXT, expires INTEGER, trusted INTEGER, PRIMARY KEY(tenant,id))')
    conn.executemany('INSERT INTO memory VALUES (?, ?, ?, ?, ?, ?)',
        ((f'tenant-{i%4}', i//4, 'x'*256, f'fixture:{i}', NOW+3600, 1) for i in range(count)))
    conn.commit()
    return conn


def security_checks(conn, path):
    assert retrieve(conn, 'tenant-0', 0)[1] == 'fixture:0'
    assert retrieve(conn, 'tenant-1', 0)[1] == 'fixture:1'
    for tenant in ('', None, 'tenant-unknown'):
        try:
            retrieve(conn, tenant, 0)
        except ValueError:
            pass
        else:
            raise AssertionError('Unbound tenant accepted')
    conn.execute("UPDATE memory SET trusted=0 WHERE tenant='tenant-0' AND id=1")
    conn.execute("UPDATE memory SET expires=? WHERE tenant='tenant-0' AND id=2", (NOW,))
    conn.execute("DELETE FROM memory WHERE tenant='tenant-0' AND id=3")
    conn.commit()
    assert all(retrieve(conn, 'tenant-0', i) is None for i in (1, 2, 3))
    assert retrieve(conn, 'tenant-1', 3) is not None
    export = conn.execute('SELECT tenant,id,body,source,expires,trusted FROM memory WHERE tenant=?', ('tenant-0',)).fetchall()
    assert export and all(row[0] == 'tenant-0' and row[3].startswith('fixture:') for row in export)
    assert all(row[1] != 3 for row in export)
    conn.execute('BEGIN IMMEDIATE')
    other = sqlite3.connect(path, timeout=0)
    try:
        other.execute("UPDATE memory SET body='fixture' WHERE tenant='tenant-1' AND id=0")
    except sqlite3.OperationalError as error:
        assert 'locked' in str(error)
    else:
        raise AssertionError('Expected competing writer failure')
    finally:
        other.close()
        conn.rollback()
    return ['tenant_key_isolation', 'missing_unknown_tenant_rejected', 'untrusted_excluded',
            'expired_excluded', 'delete_and_export', 'other_tenant_preserved', 'competing_writer_locked']


def run_reader(path, count, worker, iterations=500):
    conn = sqlite3.connect(path)
    samples = []
    # Warm up each connection; exclude warmup and connect time from query latency.
    for _ in range(20):
        retrieve(conn, f'tenant-{worker%4}', count//4-1)
    for i in range(iterations):
        start = time.perf_counter_ns()
        row = retrieve(conn, f'tenant-{worker%4}', 10 + i % (count//4-10))
        samples.append((time.perf_counter_ns()-start)/1_000_000)
        assert row is not None
    conn.close()
    return samples


def percentile(samples, p):
    return round(sorted(samples)[max(0, math.ceil(len(samples)*p)-1)], 4)


def main():
    output = {'environment': {'python': platform.python_version(), 'sqlite': sqlite3.sqlite_version,
             'system': platform.system(), 'release': platform.release(), 'machine': platform.machine(),
             'logical_cpus': os.cpu_count()}, 'method': '256-byte body; exact-key reads; 4 tenants; warm connections; no network, embedding or model calls', 'runs': []}
    with tempfile.TemporaryDirectory(prefix='memory-reference-') as directory:
        for count in (1000, 10000):
            path = Path(directory) / f'memory-{count}.sqlite'
            conn = seed(path, count)
            checks = security_checks(conn, path)
            for concurrency in (1, 4):
                started = time.perf_counter()
                with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool:
                    parts = list(pool.map(lambda w: run_reader(path, count, w), range(concurrency)))
                elapsed = time.perf_counter()-started
                samples = [sample for part in parts for sample in part]
                output['runs'].append({'initial_rows': count, 'readers': concurrency, 'queries': len(samples),
                    'p50_ms': percentile(samples, .5), 'p95_ms': percentile(samples, .95),
                    'p99_ms': percentile(samples, .99), 'wall_seconds': round(elapsed, 4),
                    'queries_per_second': round(len(samples)/elapsed), 'read_errors': 0})
            conn.execute('PRAGMA wal_checkpoint(TRUNCATE)')
            conn.close()
            output.setdefault('storage', []).append({'initial_rows': count, 'database_bytes': path.stat().st_size})
        output['checks_passed'] = checks
    output['cleanup'] = 'Temporary databases removed automatically; logical deletion is not media erasure.'
    print(json.dumps(output, indent=2))


if __name__ == '__main__':
    main()