/*
 * Define the Protocol Buffer schema for our Bio-Engineering Vector Search.
 * Save this content as `bio_vector_search.proto`
 */
syntax = "proto3";

package bio_vector_search.v1;

// Represents a biological embedding (e.g., molecular, protein)
message Embedding {
  string id = 1; // Unique identifier (e.g., PDB ID, SMILES string, Gene ID)
  repeated float vector = 2; // The embedding vector
  map<string, string> metadata = 3; // Optional metadata (e.g., source, type)
}

// Request for a vector search
message SearchRequest {
  Embedding query_embedding = 1;
  int32 k = 2; // Number of top similar items to return
  map<string, string> filter_criteria = 3; // Optional filtering based on metadata
}

// A single search result item
message SearchResultItem {
  Embedding result_embedding = 1; // The embedding found
  float similarity_score = 2; // Similarity score (e.g., cosine similarity)
}

// Response containing search results
message SearchResponse {
  repeated SearchResultItem results = 1;
}

// The service definition for vector search
service VectorSearchService {
  // Performs a similarity search for a given embedding
  rpc Search (SearchRequest) returns (SearchResponse) {}
  // Add more RPCs as needed, e.g., for batch search, index update
}

/*
 * Command-line instructions to compile the `bio_vector_search.proto` file
 * into Java and Python source code. Execute these in your terminal.
 *
 * Prerequisites: Ensure `protoc` (Protocol Buffer Compiler) is installed
 * and in your PATH. Install gRPC plugins for Java and Python.
 *
 * For Java: Add the `grpc-netty-shaded` and `protoc-gen-grpc-java` dependencies
 * to your `pom.xml` or `build.gradle`.
 * For Python: Install `grpcio` and `grpcio-tools` via pip.
 */

# Compile for Java (assuming `bio_vector_search.proto` is in current directory):
protoc --plugin=protoc-gen-grpc=/path/to/protoc-gen-grpc-java \ 
       --java_out=java/src/main/java \ 
       --grpc_out=java/src/main/java \ 
       bio_vector_search.proto

# Note: Adjust `/path/to/protoc-gen-grpc-java` to your actual plugin path.
# If using Maven/Gradle, this is often handled automatically by plugins.

# Compile for Python:
python -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. bio_vector_search.proto

# This generates:
# - `bio_vector_search_pb2.py` (message definitions)
# - `bio_vector_search_pb2_grpc.py` (service stubs and servants)

/*
 * Java gRPC Client Implementation Example.
 * This code demonstrates how a Java microservice would initiate a search request
 * to the Python vector search service.
 *
 * Prerequisites: Ensure you have compiled the `bio_vector_search.proto`
 * into Java sources as described in Part 2 and included them in your project.
 * Also, ensure gRPC dependencies are in your `pom.xml` or `build.gradle`.
 */

package com.bio.search.client;

import bio_vector_search.v1.BioVectorSearch.*; // Generated messages and services
import bio_vector_search.v1.VectorSearchServiceGrpc;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.StatusRuntimeException;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

public class BioVectorSearchClient {

    private final ManagedChannel channel;
    private final VectorSearchServiceGrpc.VectorSearchServiceBlockingStub blockingStub;

    /** Construct client for accessing VectorSearchService server. */
    public BioVectorSearchClient(String host, int port) {
        this(ManagedChannelBuilder.forAddress(host, port)
                .usePlaintext() // Use `useTransportSecurity()` for production with TLS
                .build());
    }

    BioVectorSearchClient(ManagedChannel channel) {
        this.channel = channel;
        blockingStub = VectorSearchServiceGrpc.newBlockingStub(channel);
    }

    public void shutdown() throws InterruptedException {
        channel.shutdown().awaitTermination(5, TimeUnit.SECONDS);
    }

    /** Searches for similar embeddings. */
    public void searchEmbeddings(String queryId, float[] queryVector, int k) {
        System.out.println("*** Searching for embedding: " + queryId + " ***");

        Embedding queryEmbedding = Embedding.newBuilder()
                .setId(queryId)
                .addAllVector(Arrays.asList(toObject(queryVector)))
                .putMetadata("source", "JavaClient")
                .build();

        SearchRequest request = SearchRequest.newBuilder()
                .setQueryEmbedding(queryEmbedding)
                .setK(k)
                .putFilterCriteria("type", "protein") // Example filter
                .build();

        SearchResponse response;
        try {
            response = blockingStub.search(request);
            System.out.println("Search successful. Found " + response.getResultsCount() + " results.");
            for (SearchResultItem item : response.getResultsList()) {
                System.out.println("  - ID: " + item.getResultEmbedding().getId() +
                                   ", Score: " + item.getSimilarityScore() +
                                   ", Metadata: " + item.getResultEmbedding().getMetadataMap());
            }
        } catch (StatusRuntimeException e) {
            System.err.println("RPC failed: " + e.getStatus());
            return;
        }
    }

    // Helper to convert primitive float array to Float object list for Protobuf
    private static Float[] toObject(float[] array) {
        Float[] oArray = new Float[array.length];
        for (int i = 0; i < array.length; i++) {
            oArray[i] = array[i];
        }
        return oArray;
    }

    public static void main(String[] args) throws Exception {
        BioVectorSearchClient client = new BioVectorSearchClient("localhost", 50051);
        try {
            // Example query embedding
            float[] exampleVector = {0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f, 0.9f, 1.0f};
            client.searchEmbeddings("ProteinX_123", exampleVector, 5);

            float[] anotherVector = {0.9f, 0.8f, 0.7f, 0.6f, 0.5f, 0.4f, 0.3f, 0.2f, 0.1f, 0.0f};
            client.searchEmbeddings("MoleculeY_456", anotherVector, 3);

        } finally {
            client.shutdown();
        }
    }
}

/*
 * Python gRPC Server Implementation Example.
 * This code demonstrates how a Python service can host a vector search index
 * and respond to gRPC requests from Java microservices.
 *
 * Prerequisites: Ensure you have compiled the `bio_vector_search.proto`
 * into Python sources as described in Part 2 and included them in your project.
 * Install `grpcio` and `grpcio-tools` via pip.
 */

import grpc
import time
import numpy as np
from concurrent import futures

# Import generated protobuf classes
import bio_vector_search_pb2 as pb2
import bio_vector_search_pb2_grpc as pb2_grpc

_ONE_DAY_IN_SECONDS = 60 * 60 * 24

# --- Dummy Vector Index (Replace with actual FAISS/Annoy/HNSWlib) ---
class DummyVectorIndex:
    def __init__(self, dimensions):
        self.dimensions = dimensions
        self.embeddings = [] # Stores pb2.Embedding objects
        self.vectors = [] # Stores numpy arrays
        self.id_to_index = {}
        self._populate_dummy_data()

    def _populate_dummy_data(self):
        # Simulate a pre-built index with some known embeddings
        for i in range(100):
            vec = np.random.rand(self.dimensions).astype(np.float32)
            emb_id = f"MOL_{i:03d}"
            embedding = pb2.Embedding(id=emb_id, vector=vec.tolist(), metadata={"type": "molecule"})
            self.add_embedding(embedding)

        # Add a specific embedding that's 'close' to one of our query examples
        close_vec = np.array([0.05, 0.15, 0.25, 0.35, 0.45, 0.55, 0.65, 0.75, 0.85, 0.95]).astype(np.float32)
        close_emb = pb2.Embedding(id="PROT_CLOSE", vector=close_vec.tolist(), metadata={"type": "protein"})
        self.add_embedding(close_emb)

    def add_embedding(self, embedding: pb2.Embedding):
        if embedding.id not in self.id_to_index:
            self.id_to_index[embedding.id] = len(self.embeddings)
            self.embeddings.append(embedding)
            self.vectors.append(np.array(embedding.vector, dtype=np.float32))
        else:
            print(f"Warning: Embedding ID {embedding.id} already exists. Skipping.")

    def search(self, query_vector: np.ndarray, k: int, filter_criteria: dict = None):
        if not self.vectors:
            return []

        # In a real system, use an optimized library like FAISS/Annoy/HNSWlib
        # For dummy: calculate cosine similarity with all vectors
        query_norm = np.linalg.norm(query_vector)
        if query_norm == 0:
            return [] # Avoid division by zero

        scores = []
        for i, vec in enumerate(self.vectors):
            # Apply filters first (dummy implementation)
            if filter_criteria:
                match = True
                for key, value in filter_criteria.items():
                    if item.metadata.get(key) != value:
                        match = False
                        break
                if not match:
                    continue

            vec_norm = np.linalg.norm(vec)
            if vec_norm == 0:
                similarity = 0.0
            else:
                similarity = np.dot(query_vector, vec) / (query_norm * vec_norm)
            scores.append((similarity, self.embeddings[i]))
        
        # Sort by similarity in descending order
        scores.sort(key=lambda x: x[0], reverse=True)

        # Return top K, applying filters if any in a more robust way
        results = []
        for score, embedding in scores:
            if len(results) >= k:
                break
            # Re-apply filter for actual result count logic if needed for the dummy search
            if filter_criteria:
                item_match = True
                for key, value in filter_criteria.items():
                    if embedding.metadata.get(key) != value:
                        item_match = False
                        break
                if not item_match: # Skip if it doesn't match filter
                    continue
            
            results.append(pb2.SearchResultItem(
                result_embedding=embedding,
                similarity_score=score
            ))
        return results

# --- gRPC Servicer Implementation ---
class VectorSearchServicer(pb2_grpc.VectorSearchServiceServicer):
    def __init__(self, index: DummyVectorIndex):
        self.index = index

    def Search(self, request: pb2.SearchRequest, context) -> pb2.SearchResponse:
        print(f"Received search request for ID: {request.query_embedding.id}, k={request.k}, filters={request.filter_criteria}")
        query_vector = np.array(request.query_embedding.vector, dtype=np.float32)
        
        # Perform the search using the underlying index
        search_results = self.index.search(query_vector, request.k, request.filter_criteria)

        # Build the response
        return pb2.SearchResponse(results=search_results)


def serve():
    # Initialize dummy index with 10 dimensions, matching our Java example
    dummy_index = DummyVectorIndex(dimensions=10)

    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    pb2_grpc.add_VectorSearchServiceServicer_to_server(
        VectorSearchServicer(dummy_index), server)
    server.add_insecure_port('[::]:50051') # Use `add_secure_port` for production with TLS
    server.start()
    print("Python gRPC Server started on port 50051.")
    try:
        while True:
            time.sleep(_ONE_DAY_IN_SECONDS)
    except KeyboardInterrupt:
        server.stop(0)
        print("Python gRPC Server stopped.")


if __name__ == '__main__':
    serve()

/*
 * Illustrative Python gRPC Interceptor (concept for security/logging).
 * This code snippet demonstrates the concept of a server-side interceptor
 * to add cross-cutting concerns like authentication or logging.
 *
 * In a production scenario, you'd integrate this with a proper authentication
 * system (e.g., JWT validation) and add it to your gRPC server setup.
 */

import grpc
from concurrent import futures
import time

# Assuming bio_vector_search_pb2_grpc and VectorSearchServicer are defined as before
# from your generated files and custom servicer.

# Placeholder for a simple authentication logic
def authenticate_request(context):
    metadata = dict(context.invocation_metadata())
    auth_token = metadata.get('authorization')
    if auth_token == 'Bearer_ValidToken123':
        print("Authentication successful.")
        return True
    print("Authentication failed: Invalid or missing token.")
    context.abort(grpc.StatusCode.UNAUTHENTICATED, 'Invalid or missing authentication token.')
    return False

class AuthInterceptor(grpc.ServerInterceptor):
    def intercept_service(self, continuation, handler_call_details):
        # Example: Intercept all RPCs
        method = handler_call_details.method
        # Only intercept methods that require authentication
        if '/bio_vector_search.v1.VectorSearchService/Search' in method:
            # Get the context from handler_call_details to pass to authenticate_request
            # This requires a bit more advanced setup, typically done through `grpc.ServicerContext`
            # For simplicity, we illustrate the concept here. In real interceptors,
            # you'd extract `context` correctly from the handler object.
            # For this example, let's assume `authenticate_request` gets what it needs.
            
            # NOTE: Directly accessing `context` inside interceptor requires specific gRPC patterns.
            # For a quick demo, we'll bypass the `context` argument for `authenticate_request`
            # and focus on the conceptual flow.

            # Conceptual call to authentication
            # If authenticate_request fails, it calls context.abort(), which raises an exception.
            # A real implementation would pass `handler_call_details.invocation_metadata()`
            # and potentially a custom context object.
            
            # For demo, we'll simulate authentication state.
            # If you want to run this, you need to manually add an 'authorization' header
            # to your Java client metadata to pass this check.
            print(f"Intercepting method: {method}")
            # Simulate an authentication check. In a real scenario, `authenticate_request`
            # would use the `context` from the handler to abort the call if unauthenticated.
            # This simplified demo assumes an external authentication mechanism for brevity.
            # If you remove the `authenticate_request` call, this interceptor will just log.
            # authenticate_request(handler_call_details)

        return continuation(handler_call_details)

def serve_with_interceptors():
    dummy_index = DummyVectorIndex(dimensions=10) # Re-use from Part 4
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10), interceptors=[AuthInterceptor()])
    pb2_grpc.add_VectorSearchServiceServicer_to_server(
        VectorSearchServicer(dummy_index), server)
    
    # Secure port for production: requires SSL/TLS certificates
    # with open('server.key', 'rb') as f: private_key = f.read()
    # with open('server.pem', 'rb') as f: certificate_chain = f.read()
    # server_credentials = grpc.ssl_server_credentials(((private_key, certificate_chain),))
    # server.add_secure_port('[::]:50051', server_credentials)

    server.add_insecure_port('[::]:50051') # For demo purposes
    server.start()
    print("Python gRPC Server (with interceptors) started on port 50051.")
    try:
        while True:
            time.sleep(_ONE_DAY_IN_SECONDS)
    except KeyboardInterrupt:
        server.stop(0)
        print("Python gRPC Server stopped.")

if __name__ == '__main__':
    serve_with_interceptors()

/*
 * To test the Java client with an Authorization header, modify the client:
 */
// In BioVectorSearchClient.java, modify the searchEmbeddings method:
// ...
// SearchRequest request = SearchRequest.newBuilder()
//         .setQueryEmbedding(queryEmbedding)
//         .setK(k)
//         .putFilterCriteria("type", "protein")
//         .build();

// Add a call option with metadata (e.g., Authorization header)
// You would typically use ClientInterceptor for this in a real app,
// but for demonstration, direct attachment:
// Metadata headers = new Metadata();
// headers.put(Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER), "Bearer_ValidToken123");
// response = blockingStub.withMetadata(headers).search(request);
// ...