#!/usr/bin/env python3
"""
vector_search.py - SIMD-accelerated vector search sidecar for CIE-10 SISMaule.

Listens on a Unix domain socket. For each request:
  - Receives: 4096 bytes (1024 float32 query vector)
  - Returns:  80 bytes (20 uint32 indices, descending score order)

The index matrix (12810 x 1024 float32) is loaded once into RAM at startup.
numpy uses OpenBLAS with AVX2+FMA for the matrix-vector multiply (~5-10ms).

Socket: /tmp/cie10_search.sock
Index:  /var/www/cie10/cie10_index.bin (same binary format as PHP reader)
        Header: 8 bytes (uint32 count, uint32 dims), then raw float32 data.

Usage:
    python3 /var/www/cie10/vector_search.py &
    # or via systemd (recommended)

Restart policy: if index file is not found, exits with code 1.
"""

import os
import sys
import struct
import socket
import logging
import numpy as np

# ── Config ────────────────────────────────────────────────────────────────────
SOCKET_PATH = '/var/www/cie10/vector_search.sock'
INDEX_PATH  = '/var/www/cie10/cie10_index.bin'
TOP_K       = 20
DIMS        = 1024
LOG_LEVEL   = logging.INFO

# ── Logging ───────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=LOG_LEVEL,
    format='%(asctime)s [vector_search] %(levelname)s %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
)
log = logging.getLogger(__name__)

# ── Load index ────────────────────────────────────────────────────────────────
log.info(f'Loading index from {INDEX_PATH}...')

if not os.path.isfile(INDEX_PATH):
    log.error(f'Index file not found: {INDEX_PATH}')
    sys.exit(1)

with open(INDEX_PATH, 'rb') as f:
    header = f.read(8)
    count, dims = struct.unpack('<II', header)
    raw = f.read()

log.info(f'Index header: {count} entries x {dims} dims')

if dims != DIMS:
    log.error(f'Unexpected dims: got {dims}, expected {DIMS}')
    sys.exit(1)

# Load into a contiguous C-order float32 matrix — optimal for BLAS matmul
INDEX = np.frombuffer(raw, dtype=np.float32).reshape(count, dims).copy()
log.info(f'Index loaded: {INDEX.shape}, {INDEX.nbytes / 1024 / 1024:.1f} MB')
log.info(f"numpy {np.__version__} — BLAS: OK")

# ── Socket setup ──────────────────────────────────────────────────────────────
if os.path.exists(SOCKET_PATH):
    os.unlink(SOCKET_PATH)

server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
server.bind(SOCKET_PATH)
os.chmod(SOCKET_PATH, 0o666)  # allow www-data to connect
server.listen(32)

log.info(f'Listening on {SOCKET_PATH}')

QUERY_BYTES = DIMS * 4        # 1024 floats * 4 bytes = 4096 bytes
RESULT_BYTES = TOP_K * 4      # 20 uint32 * 4 bytes = 80 bytes

# ── Request loop ──────────────────────────────────────────────────────────────
def recv_exact(conn: socket.socket, n: int) -> bytes:
    """Read exactly n bytes from socket, handling partial reads."""
    buf = bytearray()
    while len(buf) < n:
        chunk = conn.recv(n - len(buf))
        if not chunk:
            raise ConnectionError('client disconnected')
        buf.extend(chunk)
    return bytes(buf)

log.info('Ready — waiting for requests')

while True:
    try:
        conn, _ = server.accept()
        try:
            # Receive query vector
            raw_query = recv_exact(conn, QUERY_BYTES)
            query = np.frombuffer(raw_query, dtype=np.float32)

            # SIMD matrix-vector multiply (AVX2 + FMA via OpenBLAS)
            # This is the entire vector search — ~5-10ms for 12810 x 1024
            scores = INDEX @ query

            # Top-K indices (argpartition is O(n), faster than full argsort)
            top_indices = np.argpartition(scores, -TOP_K)[-TOP_K:]
            top_indices = top_indices[np.argsort(scores[top_indices])[::-1]]

            # Send result: 20 uint32 indices
            conn.sendall(struct.pack(f'{TOP_K}I', *top_indices.astype(np.uint32)))

        except Exception as e:
            log.warning(f'Request error: {e}')
        finally:
            conn.close()

    except KeyboardInterrupt:
        log.info('Shutting down')
        server.close()
        if os.path.exists(SOCKET_PATH):
            os.unlink(SOCKET_PATH)
        sys.exit(0)
    except Exception as e:
        log.error(f'Accept error: {e}')