"""
DiskANN-specific embedding server
"""

import argparse
import json
import logging
import os
import sys
import threading
import time
from pathlib import Path
from typing import Optional

import numpy as np
import zmq

# Set up logging based on environment variable
LOG_LEVEL = os.getenv("LEANN_LOG_LEVEL", "WARNING").upper()
logger = logging.getLogger(__name__)

# Force set logger level (don't rely on basicConfig in subprocess)
log_level = getattr(logging, LOG_LEVEL, logging.WARNING)
logger.setLevel(log_level)

# Ensure we have a handler if none exists
if not logger.handlers:
    handler = logging.StreamHandler()
    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.propagate = False


def create_diskann_embedding_server(
    passages_file: Optional[str] = None,
    zmq_port: int = 5555,
    model_name: str = "sentence-transformers/all-mpnet-base-v2",
    embedding_mode: str = "sentence-transformers",
    distance_metric: str = "l2",
):
    """
    Create and start a ZMQ-based embedding server for DiskANN backend.
    Uses ROUTER socket and protobuf communication as required by DiskANN C++ implementation.
    """
    logger.info(f"Starting DiskANN server on port {zmq_port} with model {model_name}")
    logger.info(f"Using embedding mode: {embedding_mode}")

    # Add leann-core to path for unified embedding computation
    current_dir = Path(__file__).parent
    leann_core_path = current_dir.parent.parent / "leann-core" / "src"
    sys.path.insert(0, str(leann_core_path))

    try:
        from leann.api import PassageManager
        from leann.embedding_compute import compute_embeddings

        logger.info("Successfully imported unified embedding computation module")
    except ImportError as e:
        logger.error(f"Failed to import embedding computation module: {e}")
        return
    finally:
        sys.path.pop(0)

    # Check port availability
    import socket

    def check_port(port):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            return s.connect_ex(("localhost", port)) == 0

    if check_port(zmq_port):
        logger.error(f"Port {zmq_port} is already in use")
        return

    # Only support metadata file, fail fast for everything else
    if not passages_file or not passages_file.endswith(".meta.json"):
        raise ValueError("Only metadata files (.meta.json) are supported")

    # Load metadata to get passage sources
    with open(passages_file) as f:
        meta = json.load(f)

    logger.info(f"Loading PassageManager with metadata_file_path: {passages_file}")
    passages = PassageManager(meta["passage_sources"], metadata_file_path=passages_file)
    logger.info(
        f"Loaded PassageManager with {len(passages.global_offset_map)} passages from metadata"
    )

    # Import protobuf after ensuring the path is correct
    try:
        from . import embedding_pb2
    except ImportError as e:
        logger.error(f"Failed to import protobuf module: {e}")
        return

    def zmq_server_thread():
        """ZMQ server thread using REP socket for universal compatibility"""
        context = zmq.Context()
        socket = context.socket(
            zmq.REP
        )  # REP socket for both BaseSearcher and DiskANN C++ REQ clients
        socket.bind(f"tcp://*:{zmq_port}")
        logger.info(f"DiskANN ZMQ REP server listening on port {zmq_port}")

        socket.setsockopt(zmq.RCVTIMEO, 1000)
        socket.setsockopt(zmq.SNDTIMEO, 1000)
        socket.setsockopt(zmq.LINGER, 0)

        while True:
            try:
                # REP socket receives single-part messages
                message = socket.recv()

                # Check for empty messages - REP socket requires response to every request
                if len(message) == 0:
                    logger.debug("Received empty message, sending empty response")
                    socket.send(b"")  # REP socket must respond to every request
                    continue

                logger.debug(f"Received ZMQ request of size {len(message)} bytes")
                logger.debug(f"Message preview: {message[:50]}")  # Show first 50 bytes

                e2e_start = time.time()

                # Try protobuf first (for DiskANN C++ node_ids requests - primary use case)
                texts = []
                node_ids = []
                is_text_request = False

                try:
                    req_proto = embedding_pb2.NodeEmbeddingRequest()
                    req_proto.ParseFromString(message)
                    node_ids = list(req_proto.node_ids)

                    if not node_ids:
                        raise RuntimeError(
                            f"PROTOBUF: Received empty node_ids! Message size: {len(message)}"
                        )

                    logger.info(
                        f"✅ PROTOBUF: Node ID request for {len(node_ids)} node embeddings: {node_ids[:10]}"
                    )
                except Exception as protobuf_error:
                    logger.debug(f"Protobuf parsing failed: {protobuf_error}")
                    # Fallback to msgpack (for BaseSearcher direct text requests)
                    try:
                        import msgpack

                        request = msgpack.unpackb(message)
                        # For BaseSearcher compatibility, request is a list of texts directly
                        if isinstance(request, list) and all(
                            isinstance(item, str) for item in request
                        ):
                            texts = request
                            is_text_request = True
                            logger.info(f"✅ MSGPACK: Direct text request for {len(texts)} texts")
                        else:
                            raise ValueError("Not a valid msgpack text request")
                    except Exception as msgpack_error:
                        raise RuntimeError(
                            f"Both protobuf and msgpack parsing failed! Protobuf: {protobuf_error}, Msgpack: {msgpack_error}"
                        )

                # Look up texts by node IDs (only if not direct text request)
                if not is_text_request:
                    for nid in node_ids:
                        try:
                            passage_data = passages.get_passage(str(nid))
                            txt = passage_data["text"]
                            if not txt:
                                raise RuntimeError(f"FATAL: Empty text for passage ID {nid}")
                            texts.append(txt)
                        except KeyError as e:
                            logger.error(f"Passage ID {nid} not found: {e}")
                            raise e
                        except Exception as e:
                            logger.error(f"Exception looking up passage ID {nid}: {e}")
                            raise

                    # Debug logging
                    logger.debug(f"Processing {len(texts)} texts")
                    logger.debug(f"Text lengths: {[len(t) for t in texts[:5]]}")  # Show first 5

                # Process embeddings using unified computation
                embeddings = compute_embeddings(texts, model_name, mode=embedding_mode)
                logger.info(
                    f"Computed embeddings for {len(texts)} texts, shape: {embeddings.shape}"
                )

                # Prepare response based on request type
                if is_text_request:
                    # For BaseSearcher compatibility: return msgpack format
                    import msgpack

                    response_data = msgpack.packb(embeddings.tolist())
                else:
                    # For DiskANN C++ compatibility: return protobuf format
                    resp_proto = embedding_pb2.NodeEmbeddingResponse()
                    hidden_contiguous = np.ascontiguousarray(embeddings, dtype=np.float32)

                    # Serialize embeddings data
                    resp_proto.embeddings_data = hidden_contiguous.tobytes()
                    resp_proto.dimensions.append(hidden_contiguous.shape[0])
                    resp_proto.dimensions.append(hidden_contiguous.shape[1])

                    response_data = resp_proto.SerializeToString()

                # Send response back to the client
                socket.send(response_data)

                e2e_end = time.time()
                logger.info(f"⏱️  ZMQ E2E time: {e2e_end - e2e_start:.6f}s")

            except zmq.Again:
                logger.debug("ZMQ socket timeout, continuing to listen")
                continue
            except Exception as e:
                logger.error(f"Error in ZMQ server loop: {e}")
                import traceback

                traceback.print_exc()
                raise

    def zmq_server_thread_with_shutdown(shutdown_event):
        """ZMQ server thread that respects shutdown signal.

        This creates its own REP socket, binds to zmq_port, and periodically
        checks shutdown_event using recv timeouts to exit cleanly.
        """
        logger.info("DiskANN ZMQ server thread started with shutdown support")

        context = zmq.Context()
        rep_socket = context.socket(zmq.REP)
        rep_socket.bind(f"tcp://*:{zmq_port}")
        logger.info(f"DiskANN ZMQ REP server listening on port {zmq_port}")

        # Set receive timeout so we can check shutdown_event periodically
        rep_socket.setsockopt(zmq.RCVTIMEO, 1000)  # 1 second timeout
        rep_socket.setsockopt(zmq.SNDTIMEO, 1000)
        rep_socket.setsockopt(zmq.LINGER, 0)

        try:
            while not shutdown_event.is_set():
                try:
                    e2e_start = time.time()
                    # REP socket receives single-part messages
                    message = rep_socket.recv()

                    # Check for empty messages - REP socket requires response to every request
                    if not message:
                        logger.warning("Received empty message, sending empty response")
                        rep_socket.send(b"")
                        continue

                    # Try protobuf first (same logic as original)
                    texts = []
                    is_text_request = False

                    try:
                        req_proto = embedding_pb2.NodeEmbeddingRequest()
                        req_proto.ParseFromString(message)
                        node_ids = list(req_proto.node_ids)

                        # Look up texts by node IDs
                        for nid in node_ids:
                            try:
                                passage_data = passages.get_passage(str(nid))
                                txt = passage_data["text"]
                                if not txt:
                                    raise RuntimeError(f"FATAL: Empty text for passage ID {nid}")
                                texts.append(txt)
                            except KeyError:
                                raise RuntimeError(f"FATAL: Passage with ID {nid} not found")

                        logger.info(f"ZMQ received protobuf request for {len(node_ids)} node IDs")
                    except Exception:
                        # Fallback to msgpack for text requests
                        try:
                            import msgpack

                            request = msgpack.unpackb(message)
                            if isinstance(request, list) and all(
                                isinstance(item, str) for item in request
                            ):
                                texts = request
                                is_text_request = True
                                logger.info(
                                    f"ZMQ received msgpack text request for {len(texts)} texts"
                                )
                            else:
                                raise ValueError("Not a valid msgpack text request")
                        except Exception:
                            logger.error("Both protobuf and msgpack parsing failed!")
                            # Send error response
                            resp_proto = embedding_pb2.NodeEmbeddingResponse()
                            rep_socket.send(resp_proto.SerializeToString())
                            continue

                    # Process the request
                    embeddings = compute_embeddings(texts, model_name, mode=embedding_mode)
                    logger.info(f"Computed embeddings shape: {embeddings.shape}")

                    # Validation
                    if np.isnan(embeddings).any() or np.isinf(embeddings).any():
                        logger.error("NaN or Inf detected in embeddings!")
                        # Send error response
                        if is_text_request:
                            import msgpack

                            response_data = msgpack.packb([])
                        else:
                            resp_proto = embedding_pb2.NodeEmbeddingResponse()
                            response_data = resp_proto.SerializeToString()
                        rep_socket.send(response_data)
                        continue

                    # Prepare response based on request type
                    if is_text_request:
                        # For direct text requests, return msgpack
                        import msgpack

                        response_data = msgpack.packb(embeddings.tolist())
                    else:
                        # For protobuf requests, return protobuf
                        resp_proto = embedding_pb2.NodeEmbeddingResponse()
                        hidden_contiguous = np.ascontiguousarray(embeddings, dtype=np.float32)

                        resp_proto.embeddings_data = hidden_contiguous.tobytes()
                        resp_proto.dimensions.append(hidden_contiguous.shape[0])
                        resp_proto.dimensions.append(hidden_contiguous.shape[1])

                        response_data = resp_proto.SerializeToString()

                    # Send response back to the client
                    rep_socket.send(response_data)

                    e2e_end = time.time()
                    logger.info(f"⏱️  ZMQ E2E time: {e2e_end - e2e_start:.6f}s")

                except zmq.Again:
                    # Timeout - check shutdown_event and continue
                    continue
                except Exception as e:
                    if not shutdown_event.is_set():
                        logger.error(f"Error in ZMQ server loop: {e}")
                        try:
                            # Send error response for REP socket
                            resp_proto = embedding_pb2.NodeEmbeddingResponse()
                            rep_socket.send(resp_proto.SerializeToString())
                        except Exception:
                            pass
                    else:
                        logger.info("Shutdown in progress, ignoring ZMQ error")
                        break
        finally:
            try:
                rep_socket.close(0)
            except Exception:
                pass
            try:
                context.term()
            except Exception:
                pass

        logger.info("DiskANN ZMQ server thread exiting gracefully")

    # Add shutdown coordination
    shutdown_event = threading.Event()

    def shutdown_zmq_server():
        """Gracefully shutdown ZMQ server."""
        logger.info("Initiating graceful shutdown...")
        shutdown_event.set()

        if zmq_thread.is_alive():
            logger.info("Waiting for ZMQ thread to finish...")
            zmq_thread.join(timeout=5)
            if zmq_thread.is_alive():
                logger.warning("ZMQ thread did not finish in time")

        # Clean up ZMQ resources
        try:
            # Note: socket and context are cleaned up by thread exit
            logger.info("ZMQ resources cleaned up")
        except Exception as e:
            logger.warning(f"Error cleaning ZMQ resources: {e}")

        # Clean up other resources
        try:
            import gc

            gc.collect()
            logger.info("Additional resources cleaned up")
        except Exception as e:
            logger.warning(f"Error cleaning additional resources: {e}")

        logger.info("Graceful shutdown completed")
        sys.exit(0)

    # Register signal handlers within this function scope
    import signal

    def signal_handler(sig, frame):
        logger.info(f"Received signal {sig}, shutting down gracefully...")
        shutdown_zmq_server()

    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)

    # Start ZMQ thread (NOT daemon!)
    zmq_thread = threading.Thread(
        target=lambda: zmq_server_thread_with_shutdown(shutdown_event),
        daemon=False,  # Not daemon - we want to wait for it
    )
    zmq_thread.start()
    logger.info(f"Started DiskANN ZMQ server thread on port {zmq_port}")

    # Keep the main thread alive
    try:
        while not shutdown_event.is_set():
            time.sleep(0.1)  # Check shutdown more frequently
    except KeyboardInterrupt:
        logger.info("DiskANN Server shutting down...")
        shutdown_zmq_server()
        return

    # If we reach here, shutdown was triggered by signal
    logger.info("Main loop exited, process should be shutting down")


if __name__ == "__main__":
    import sys

    # Signal handlers are now registered within create_diskann_embedding_server

    parser = argparse.ArgumentParser(description="DiskANN Embedding service")
    parser.add_argument("--zmq-port", type=int, default=5555, help="ZMQ port to run on")
    parser.add_argument(
        "--passages-file",
        type=str,
        help="Metadata JSON file containing passage sources",
    )
    parser.add_argument(
        "--model-name",
        type=str,
        default="sentence-transformers/all-mpnet-base-v2",
        help="Embedding model name",
    )
    parser.add_argument(
        "--embedding-mode",
        type=str,
        default="sentence-transformers",
        choices=["sentence-transformers", "openai", "mlx", "ollama"],
        help="Embedding backend mode",
    )
    parser.add_argument(
        "--distance-metric",
        type=str,
        default="l2",
        choices=["l2", "mips", "cosine"],
        help="Distance metric for similarity computation",
    )

    args = parser.parse_args()

    # Create and start the DiskANN embedding server
    create_diskann_embedding_server(
        passages_file=args.passages_file,
        zmq_port=args.zmq_port,
        model_name=args.model_name,
        embedding_mode=args.embedding_mode,
        distance_metric=args.distance_metric,
    )
