import os
import uuid
import logging
from urllib import request
#
import spacy
from fastmcp import FastMCP
from qdrant_client import AsyncQdrantClient, models
from fastembed import TextEmbedding

# logger config
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(name=__name__)

# 1. Initialize FastMCP with your server name
mcp = FastMCP("Qdrant Cloud Tools")

# 2. Open an asynchronous network client to talk to the local containerized Qdrant Engine
qdrant_cl = AsyncQdrantClient(url="http://localhost:6333")

# 3. Spin up an ONNX embedding runner (lightweight, runs cleanly on CPU)
# This model outputs 384-dimensional dense vectors
EMBED_MODEL_NAME = "BAAI/bge-small-en-v1.5"
VECTOR_DIMENSION = 384
COLLECTION_NAME = "knowledge_base"
embedding_engine = TextEmbedding(model_name=EMBED_MODEL_NAME)

# 4. Initialize a lightweight SpaCy sentence boundary processor
nlp = spacy.blank("en")
nlp.add_pipe("sentencizer")

# Check IP address
try:
    ip = request.urlopen("https://api.ipify.org").read().decode()
    logger.info(f"IP Address: {ip}")
except Exception as e:
    logger.error(repr(e))
    

def chunk_text_by_sentence(text: str, max_chunk_chars: int = 1200, overlap_sentences: int = 1) -> list[str]:
    """
    Parses document text into clean chunks using SpaCy sentence boundaries,
    ensuring sentences are never sliced in half.
    """
    doc = nlp(text)
    sentences = [sent.text.strip() for sent in doc.sents if sent.text.strip()]
    
    chunks = []
    current_chunk = []
    current_length = 0
    
    for idx, sentence in enumerate(sentences):
        sentence_len = len(sentence)
        
        # If adding this sentence exceeds our threshold and we already have content, roll the chunk
        if current_length + sentence_len > max_chunk_chars and current_chunk:
            chunks.append(" ".join(current_chunk))
            
            # Create overlap by pulling the last N sentences from the previous chunk
            if overlap_sentences > 0 and len(current_chunk) >= overlap_sentences:
                current_chunk = current_chunk[-overlap_sentences:]
                current_length = sum(len(s) for s in current_chunk) + len(current_chunk) - 1
            else:
                current_chunk = []
                current_length = 0
                
        current_chunk.append(sentence)
        current_length += sentence_len + 1  # Account for joining spaces
        
    # Pick up any remaining dangling sentences
    if current_chunk:
        chunks.append(" ".join(current_chunk))
        
    return chunks


async def ensure_collection():
    """Utility to verify or create the vector schema layout before database operations."""
    exists = await qdrant_cl.collection_exists(COLLECTION_NAME)
    if not exists:
        await qdrant_cl.create_collection(
            collection_name=COLLECTION_NAME,
            vectors_config=models.VectorParams(
                size=VECTOR_DIMENSION,
                distance=models.Distance.COSINE
            )
        )


@mcp.tool()
async def index_document(text: str, source: str = "mcp_upload") -> str:
    """
    Semantically chunks text by sentence boundaries, vectorizes, and stores items in Qdrant.
    
    Args:
        text: The raw textual data or large document to store.
        source: Context metadata indicating origin (e.g., file name, web url).
    """
    await ensure_collection()
    
    # 1. Break text down safely using SpaCy sentence grouping
    text_chunks = chunk_text_by_sentence(text, max_chunk_chars=1200, overlap_sentences=1)
    
    if not text_chunks:
        return "No text context found to index."
        
    # 2. Vectorize all semantic chunks in a single pass
    vector_generator = embedding_engine.embed(text_chunks)
    vectors = [v.tolist() for v in vector_generator]
    
    # 3. Compile a batch payload list for high-speed upsertion
    points = []
    for i, (chunk, vector) in enumerate(zip(text_chunks, vectors)):
        point_id = str(uuid.uuid4())
        points.append(
            models.PointStruct(
                id=point_id,
                vector=vector,
                payload={
                    "text": chunk, 
                    "source": source,
                    "chunk_index": i
                }
            )
        )
    
    # 4. Perform an atomic batch database upload
    await qdrant_cl.upsert(
        collection_name=COLLECTION_NAME,
        points=points
    )
    
    return f"Processed {len(points)} semantically isolated sentence-chunks from source '{source}'."


@mcp.tool()
async def semantic_search(query: str, limit: int = 3) -> str:
    """
    Query the remote Qdrant database for semantically relevant document matches.
    
    Args:
        query: The natural language search string or phrase.
        limit: Max number of document snippets to pull back.
    """
    await ensure_collection()
    
    # Convert query text into vector array matching the index space
    query_generator = embedding_engine.embed([query])
    query_vector = list(query_generator)[0].tolist()
    
    # Query vector neighbors 
    response = await qdrant_cl.query_points(
        collection_name=COLLECTION_NAME,
        query=query_vector,
        limit=limit
    )
    
    if not response.points:
        return f"No relevant matches found for query string: '{query}'."
        
    # Build clean output for the consuming client
    matches = []
    for point in response.points:
        score = round(point.score, 4)
        doc_text = point.payload.get("text", "[Data missing]")
        doc_source = point.payload.get("source", "unknown")
        matches.append(f"  - [Score: {score}] (Source: {doc_source}): {doc_text}")
        
    return f"Top semantic hits for '{query}':\n" + "\n".join(matches)


if __name__ == "__main__":
    # Let FastMCP host the network layer bound to Hugging Face's ingress port
    mcp.run(transport="streamable-http", host="0.0.0.0", port=7860)