#!/usr/bin/env python3
"""
TurboVec Sidecar Process

Scaffolding-phase sidecar for cross-session vector memory.
Reads JSON requests from stdin (one per line), writes JSON responses to stdout.

Operations:
  - init {dim, bit_width}: Create a new index
  - add {id, vector}: Add a vector to the index
  - search {vector, k}: Return top-k hits by cosine similarity
  - save {path}: Persist index to disk
  - load {path}: Load index from disk (scaffolding: partial implementation)
  - ping: Health check
  - shutdown: Exit cleanly

Requires: pip install turbovec
"""

import sys
import json
import numpy as np
from turbovec import TurboQuantIndex


def _normalize(vector):
    """L2-normalize to a (1, dim) float32 row so TurboVec's inner-product search
    behaves as cosine similarity (matching the in-memory L3 store)."""
    a = np.asarray(vector, dtype=np.float32).reshape(1, -1)
    n = np.linalg.norm(a, axis=-1, keepdims=True)
    n[n == 0] = 1.0
    return (a / n).astype(np.float32)


class TurboVecSidecar:
    def __init__(self):
        self.index = None
        self.id_to_idx = {}  # id (str) -> index position
        self.idx_to_id = []  # index position -> id (str)
        self.dim = None
        self.bit_width = 4
        self.dirty = False   # True when vectors added since last prepare()

    def init(self, dim, bit_width):
        """Initialize a new TurboQuantIndex."""
        try:
            self.index = TurboQuantIndex(dim=dim, bit_width=bit_width)
            self.dim = dim
            self.bit_width = bit_width
            self.id_to_idx = {}
            self.idx_to_id = []
            self.dirty = False
            return {"ok": True}
        except Exception as e:
            return {"ok": False, "error": str(e)}

    def add(self, entry_id, vector):
        """Add one vector. TurboVec expects a (1, dim) float32 batch, normalized."""
        try:
            if self.index is None:
                return {"ok": False, "error": "Index not initialized"}
            self.index.add(_normalize(vector))
            self.id_to_idx[entry_id] = len(self.idx_to_id)
            self.idx_to_id.append(entry_id)
            self.dirty = True
            return {"ok": True}
        except Exception as e:
            return {"ok": False, "error": str(e)}

    def search(self, vector, k):
        """Top-k inner-product (≈cosine) search via the quantized index."""
        try:
            if self.index is None:
                return {"ok": False, "error": "Index not initialized"}
            if len(self.idx_to_id) == 0:
                return {"ok": True, "hits": []}
            if self.dirty:
                self.index.prepare()   # warm caches; required after adds
                self.dirty = False
            scores, indices = self.index.search(_normalize(vector), k)
            hits = [
                {"id": self.idx_to_id[int(i)], "score": float(s)}
                for s, i in zip(scores[0], indices[0])
            ]
            return {"ok": True, "hits": hits}
        except Exception as e:
            return {"ok": False, "error": str(e)}

    def save(self, path):
        """Persist the index + id mapping to disk."""
        try:
            if self.index is None:
                return {"ok": False, "error": "Index not initialized"}
            if self.dirty:
                self.index.prepare()
                self.dirty = False
            self.index.write(path)
            metadata = {
                "id_to_idx": self.id_to_idx,
                "idx_to_id": self.idx_to_id,
                "dim": self.dim,
                "bit_width": self.bit_width,
            }
            with open(f"{path}.meta.json", "w") as f:
                json.dump(metadata, f)
            return {"ok": True}
        except Exception as e:
            return {"ok": False, "error": str(e)}

    def load(self, path):
        """Load a previously-saved index + id mapping."""
        try:
            self.index = TurboQuantIndex.load(path)
            with open(f"{path}.meta.json") as f:
                metadata = json.load(f)
            self.id_to_idx = metadata.get("id_to_idx", {})
            self.idx_to_id = metadata.get("idx_to_id", [])
            self.dim = metadata.get("dim")
            self.bit_width = metadata.get("bit_width", 4)
            self.dirty = True   # re-prepare on first search after load
            return {"ok": True}
        except Exception as e:
            return {"ok": False, "error": str(e)}

    def ping(self):
        """Health check."""
        return {"ok": True}

    def shutdown(self):
        """Exit cleanly."""
        return {"ok": True}

    def process(self, request):
        """Dispatch a request to the appropriate handler."""
        op = request.get("op")
        if op == "init":
            return self.init(request.get("dim", 1536), request.get("bitWidth", 4))
        elif op == "add":
            return self.add(request.get("id"), request.get("vector"))
        elif op == "search":
            return self.search(request.get("vector"), request.get("k", 10))
        elif op == "save":
            return self.save(request.get("path"))
        elif op == "load":
            return self.load(request.get("path"))
        elif op == "ping":
            return self.ping()
        elif op == "shutdown":
            return self.shutdown()
        else:
            return {"ok": False, "error": f"Unknown operation: {op}"}


def main():
    sidecar = TurboVecSidecar()
    try:
        while True:
            line = sys.stdin.readline()
            if not line:
                break
            try:
                request = json.loads(line.strip())
                response = sidecar.process(request)
                sys.stdout.write(json.dumps(response) + "\n")
                sys.stdout.flush()
                if request.get("op") == "shutdown":
                    break
            except json.JSONDecodeError as e:
                sys.stdout.write(json.dumps({"ok": False, "error": f"Invalid JSON: {e}"}) + "\n")
                sys.stdout.flush()
            except Exception as e:
                sys.stderr.write(f"Error processing request: {e}\n")
                sys.stdout.write(json.dumps({"ok": False, "error": str(e)}) + "\n")
                sys.stdout.flush()
    except KeyboardInterrupt:
        pass
    except Exception as e:
        sys.stderr.write(f"Fatal error: {e}\n")
        sys.exit(1)


if __name__ == "__main__":
    main()
