import os
import httpx
from mcp.server.fastmcp import FastMCP
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
import uuid
import logging

# Configuration
QDRANT_URL = os.getenv("QDRANT_URL", "http://qdrant.openshift-gitops.svc:6333")
TEI_URL = os.getenv("TEI_URL", "http://text-embeddings.tei.svc.cluster.local:8080")
EMBEDDING_DIM = 768  # BAAI/bge-base-en-v1.5

# Initialize
mcp = FastMCP("knowledge-mcp")
qdrant = QdrantClient(url=QDRANT_URL)

def get_embedding(text: str) -> list[float]:
    """Get embedding from TEI."""
    url = f"{TEI_URL}/embed" # Adjust based on TEI version, often /v1/embeddings or /embed
    # Trying standard TEI /embed endpoint for raw lists
    try:
        response = httpx.post(url, json={"inputs": text}, timeout=10.0)
        response.raise_for_status()
        return response.json()[0]
    except Exception as e:
        # Fallback to OpenAI compatible endpoint if needed
        logging.error(f"Embedding failed: {e}")
        raise

@mcp.tool()
def create_notebook(name: str) -> str:
    """Create a new RAG notebook (Qdrant collection)."""
    clean_name = name.lower().replace(" ", "-")
    
    # Check if exists
    if qdrant.collection_exists(clean_name):
        return f"Notebook '{clean_name}' already exists."

    qdrant.create_collection(
        collection_name=clean_name,
        vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.COSINE),
    )
    return f"Notebook '{clean_name}' created successfully."

@mcp.tool()
def add_source(notebook: str, text: str, source_name: str = "manual") -> str:
    """Add text content to a notebook. Ingests, chunks, and indexes."""
    if not qdrant.collection_exists(notebook):
        return f"Error: Notebook '{notebook}' does not exist."

    # Very basic chunking for now
    chunks = [text[i:i+500] for i in range(0, len(text), 500)]
    points = []

    for chunk in chunks:
        vector = get_embedding(chunk)
        points.append(PointStruct(
            id=str(uuid.uuid4()),
            vector=vector,
            payload={"source": source_name, "text": chunk}
        ))

    qdrant.upsert(collection_name=notebook, points=points)
    return f"Added {len(points)} chunks from '{source_name}' to '{notebook}'."

@mcp.tool()
def query_notebook(notebook: str, query: str, limit: int = 5) -> str:
    """Query the notebook for relevant context."""
    if not qdrant.collection_exists(notebook):
        return f"Error: Notebook '{notebook}' does not exist."

    vector = get_embedding(query)
    hits = qdrant.search(
        collection_name=notebook,
        query_vector=vector,
        limit=limit
    )

    results = []
    for hit in hits:
        results.append(f"--- (Score: {hit.score:.2f}) ---\n{hit.payload.get('text', '')}\n")

    return "\n".join(results)

if __name__ == "__main__":
    mcp.run()