#!/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
from turbovec import TurboQuantIndex


def cosine_similarity(a, b):
    """Compute cosine similarity between two vectors."""
    dot_product = sum(x * y for x, y in zip(a, b))
    norm_a = sum(x * x for x in a) ** 0.5
    norm_b = sum(x * x for x in b) ** 0.5
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot_product / (norm_a * norm_b)


class TurboVecSidecar:
    def __init__(self):
        self.index = None
        self.id_to_idx = {}  # id (str) -> index position
        self.idx_to_id = []  # index position -> id (str)

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

    def add(self, entry_id, vector):
        """Add a vector to the index."""
        try:
            if self.index is None:
                return {"ok": False, "error": "Index not initialized"}
            idx = len(self.idx_to_id)
            self.index.add(vector)
            self.id_to_idx[entry_id] = idx
            self.idx_to_id.append(entry_id)
            return {"ok": True}
        except Exception as e:
            return {"ok": False, "error": str(e)}

    def search(self, vector, k):
        """Search for top-k nearest neighbors."""
        try:
            if self.index is None:
                return {"ok": False, "error": "Index not initialized"}

            # Compute distances to all indexed vectors using cosine similarity
            hits = []
            for idx, stored_vector in enumerate(self.index.vectors):
                score = cosine_similarity(vector, stored_vector)
                entry_id = self.idx_to_id[idx]
                hits.append({"id": entry_id, "score": score})

            # Sort by score descending, take top k
            hits = sorted(hits, key=lambda x: x["score"], reverse=True)[:k]
            return {"ok": True, "hits": hits}
        except Exception as e:
            return {"ok": False, "error": str(e)}

    def save(self, path):
        """Persist the index to disk."""
        try:
            if self.index is None:
                return {"ok": False, "error": "Index not initialized"}
            self.index.write(path)
            # Also save metadata
            metadata = {"id_to_idx": self.id_to_idx, "idx_to_id": self.idx_to_id}
            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 the index from disk."""
        try:
            # SCAFFOLDING: full load not yet implemented
            # For now, return success but do not actually 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()
