#!/usr/bin/env python3
"""
Commonwealth Qdrant + Ollama Connector — mini MCP for vector search.

Usage:
    python3 qdrant_connector.py health
    python3 qdrant_connector.py create-collection
    python3 qdrant_connector.py search "some query text"
    python3 qdrant_connector.py index <id> "text to embed" [--metadata key=value ...]
    python3 qdrant_connector.py index-file <filepath>
    python3 qdrant_connector.py count
    python3 qdrant_connector.py --help

Requires only Python 3 stdlib (no pip dependencies).
"""

import hashlib
import json
import os
import re
import sys
import urllib.error
import urllib.request
from pathlib import Path

# ---------------------------------------------------------------------------
# Paths / Config
# ---------------------------------------------------------------------------

SCRIPT_DIR = Path(__file__).resolve().parent
CONFIG_PATH = SCRIPT_DIR / "config.json"


def load_config():
    """Load config.json."""
    with open(CONFIG_PATH, "r") as f:
        return json.load(f)


# ---------------------------------------------------------------------------
# HTTP helpers (stdlib only)
# ---------------------------------------------------------------------------

def http_request(url, method="GET", data=None, headers=None, timeout=30):
    """
    Perform an HTTP request using urllib.  Returns (status_code, parsed_json | raw_text).
    """
    hdrs = {"Content-Type": "application/json"}
    if headers:
        hdrs.update(headers)

    body = None
    if data is not None:
        body = json.dumps(data).encode("utf-8")

    req = urllib.request.Request(url, data=body, headers=hdrs, method=method)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8")
            try:
                return resp.status, json.loads(raw)
            except json.JSONDecodeError:
                return resp.status, raw
    except urllib.error.HTTPError as exc:
        raw = exc.read().decode("utf-8") if exc.fp else ""
        try:
            return exc.code, json.loads(raw)
        except json.JSONDecodeError:
            return exc.code, raw
    except urllib.error.URLError as exc:
        raise ConnectionError(f"Cannot reach {url}: {exc.reason}") from exc


# ---------------------------------------------------------------------------
# Embedding helper
# ---------------------------------------------------------------------------

def embed_text(cfg, text):
    """
    Call ollama /api/embed to get an embedding vector for the given text.
    Returns a list of floats.
    """
    url = f"{cfg['ollama_url']}/api/embed"
    payload = {"model": cfg["embed_model"], "input": text}
    status, resp = http_request(url, method="POST", data=payload)
    if status != 200:
        raise RuntimeError(f"Ollama embed failed (HTTP {status}): {resp}")
    # ollama returns {"embeddings": [[...]]}
    embeddings = resp.get("embeddings")
    if not embeddings or not embeddings[0]:
        raise RuntimeError(f"Ollama returned empty embeddings: {resp}")
    return embeddings[0]


# ---------------------------------------------------------------------------
# Qdrant helpers
# ---------------------------------------------------------------------------

def qdrant_create_collection(cfg):
    """Create (or recreate) the Qdrant collection."""
    url = f"{cfg['qdrant_url']}/collections/{cfg['collection']}"
    payload = {
        "vectors": {
            "size": cfg["embed_dimensions"],
            "distance": "Cosine",
        }
    }
    status, resp = http_request(url, method="PUT", data=payload)
    return status, resp


def qdrant_upsert(cfg, points):
    """Upsert a list of points into Qdrant."""
    url = f"{cfg['qdrant_url']}/collections/{cfg['collection']}/points"
    payload = {"points": points}
    status, resp = http_request(url, method="PUT", data=payload)
    return status, resp


def qdrant_search(cfg, vector, limit=5):
    """Search Qdrant by vector."""
    url = f"{cfg['qdrant_url']}/collections/{cfg['collection']}/points/query"
    payload = {"query": vector, "limit": limit, "with_payload": True}
    status, resp = http_request(url, method="POST", data=payload)
    return status, resp


def qdrant_collection_info(cfg):
    """Get collection info (includes point count)."""
    url = f"{cfg['qdrant_url']}/collections/{cfg['collection']}"
    status, resp = http_request(url, method="GET")
    return status, resp


# ---------------------------------------------------------------------------
# Chunking helper
# ---------------------------------------------------------------------------

def chunk_markdown(text, source_file=""):
    """
    Split markdown by headings (# or ##).  Returns a list of dicts:
        {"heading": str, "text": str, "chunk_index": int, "source_file": str}
    """
    # Split on lines that start with one or two hashes
    pattern = re.compile(r"^(#{1,2})\s+(.+)$", re.MULTILINE)
    matches = list(pattern.finditer(text))

    chunks = []

    if not matches:
        # No headings — treat entire file as one chunk
        stripped = text.strip()
        if stripped:
            chunks.append({
                "heading": Path(source_file).stem if source_file else "untitled",
                "text": stripped,
                "chunk_index": 0,
                "source_file": source_file,
            })
        return chunks

    # Text before the first heading
    preamble = text[: matches[0].start()].strip()
    if preamble:
        chunks.append({
            "heading": "(preamble)",
            "text": preamble,
            "chunk_index": 0,
            "source_file": source_file,
        })

    for i, match in enumerate(matches):
        heading = match.group(2).strip()
        start = match.end()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
        body = text[start:end].strip()
        if body:
            chunks.append({
                "heading": heading,
                "text": body,
                "chunk_index": len(chunks),
                "source_file": source_file,
            })

    return chunks


def text_to_point_id(text):
    """Deterministic integer ID from a string (unsigned 64-bit range for Qdrant)."""
    h = hashlib.sha256(text.encode("utf-8")).hexdigest()
    # Qdrant accepts unsigned 64-bit integer IDs
    return int(h[:16], 16)


# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------

def cmd_health(cfg):
    """Check connectivity to Qdrant and Ollama."""
    results = {}

    # Qdrant health
    try:
        status, resp = http_request(f"{cfg['qdrant_url']}/healthz", method="GET", timeout=5)
        results["qdrant"] = {"reachable": True, "status": status, "response": resp}
    except ConnectionError as exc:
        results["qdrant"] = {"reachable": False, "error": str(exc)}

    # Ollama health
    try:
        status, resp = http_request(f"{cfg['ollama_url']}/api/tags", method="GET", timeout=5)
        results["ollama"] = {"reachable": True, "status": status}
        # List available models for convenience
        if isinstance(resp, dict) and "models" in resp:
            results["ollama"]["models"] = [m.get("name", "?") for m in resp["models"]]
    except ConnectionError as exc:
        results["ollama"] = {"reachable": False, "error": str(exc)}

    all_ok = all(v.get("reachable", False) for v in results.values())
    return {"ok": all_ok, "services": results}


def cmd_create_collection(cfg):
    """Create the Qdrant collection."""
    try:
        status, resp = qdrant_create_collection(cfg)
        success = status in (200, 201)
        return {"ok": success, "status": status, "response": resp}
    except ConnectionError as exc:
        return {"ok": False, "error": str(exc)}


def cmd_search(cfg, query_text):
    """Embed query text and search Qdrant."""
    try:
        vector = embed_text(cfg, query_text)
        status, resp = qdrant_search(cfg, vector)
        if status != 200:
            return {"ok": False, "status": status, "error": resp}

        # Extract the points from the response
        points = resp.get("result", {}).get("points", resp.get("result", []))
        results = []
        if isinstance(points, list):
            for pt in points:
                results.append({
                    "id": pt.get("id"),
                    "score": pt.get("score"),
                    "payload": pt.get("payload", {}),
                })
        return {"ok": True, "query": query_text, "count": len(results), "results": results}
    except (ConnectionError, RuntimeError) as exc:
        return {"ok": False, "error": str(exc)}


def cmd_index(cfg, point_id_str, text, metadata=None):
    """Embed text and upsert a single point."""
    try:
        vector = embed_text(cfg, text)

        # Build a numeric ID from the provided string
        try:
            point_id = int(point_id_str)
        except ValueError:
            point_id = text_to_point_id(point_id_str)

        payload = metadata or {}
        payload["text"] = text

        point = {"id": point_id, "vector": vector, "payload": payload}
        status, resp = qdrant_upsert(cfg, [point])
        success = status in (200, 201)
        return {"ok": success, "status": status, "point_id": point_id, "response": resp}
    except (ConnectionError, RuntimeError) as exc:
        return {"ok": False, "error": str(exc)}


def cmd_index_file(cfg, filepath):
    """Read a markdown file, chunk it, embed each chunk, and upsert all to Qdrant."""
    fpath = Path(filepath).resolve()
    if not fpath.exists():
        return {"ok": False, "error": f"File not found: {fpath}"}

    text = fpath.read_text(encoding="utf-8")
    source = str(fpath)
    chunks = chunk_markdown(text, source_file=source)

    if not chunks:
        return {"ok": False, "error": "No content chunks extracted from file"}

    points = []
    errors = []
    for chunk in chunks:
        chunk_key = f"{source}::{chunk['heading']}::{chunk['chunk_index']}"
        point_id = text_to_point_id(chunk_key)
        try:
            vector = embed_text(cfg, chunk["text"])
        except (ConnectionError, RuntimeError) as exc:
            errors.append({"chunk": chunk["heading"], "error": str(exc)})
            continue

        points.append({
            "id": point_id,
            "vector": vector,
            "payload": {
                "source_file": chunk["source_file"],
                "heading": chunk["heading"],
                "chunk_index": chunk["chunk_index"],
                "text": chunk["text"],
            },
        })

    if not points:
        return {"ok": False, "error": "All chunks failed to embed", "details": errors}

    try:
        status, resp = qdrant_upsert(cfg, points)
        success = status in (200, 201)
        result = {
            "ok": success,
            "status": status,
            "file": source,
            "chunks_indexed": len(points),
            "chunks_failed": len(errors),
            "response": resp,
        }
        if errors:
            result["errors"] = errors
        return result
    except ConnectionError as exc:
        return {"ok": False, "error": str(exc)}


def cmd_count(cfg):
    """Return the point count in the collection."""
    try:
        status, resp = qdrant_collection_info(cfg)
        if status != 200:
            return {"ok": False, "status": status, "error": resp}
        # Qdrant returns {"result": {"points_count": N, ...}}
        result_data = resp.get("result", {})
        count = result_data.get("points_count", result_data.get("vectors_count", "unknown"))
        return {"ok": True, "collection": cfg["collection"], "points_count": count}
    except ConnectionError as exc:
        return {"ok": False, "error": str(exc)}


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

HELP_TEXT = """\
Commonwealth Qdrant + Ollama Connector

Usage:
  qdrant_connector.py health                              Check Qdrant & Ollama connectivity
  qdrant_connector.py create-collection                   Create the vector collection
  qdrant_connector.py search "<query text>"               Embed query and search Qdrant
  qdrant_connector.py index <id> "<text>" [--metadata k=v ...]
                                                          Embed text and upsert one point
  qdrant_connector.py index-file <filepath>               Chunk a markdown file and index all chunks
  qdrant_connector.py count                               Show point count in collection
  qdrant_connector.py --help                              Show this help message

All output is JSON on stdout.  Uses only Python stdlib (no pip install needed).

Config:  {config}
""".format(config=CONFIG_PATH)


def parse_metadata(args):
    """Parse --metadata key=value pairs from argument list."""
    metadata = {}
    i = 0
    while i < len(args):
        if args[i] == "--metadata" and i + 1 < len(args):
            i += 1
            while i < len(args) and "=" in args[i] and not args[i].startswith("--"):
                key, _, value = args[i].partition("=")
                metadata[key] = value
                i += 1
        else:
            i += 1
    return metadata


def main():
    if len(sys.argv) < 2 or sys.argv[1] in ("--help", "-h", "help"):
        print(HELP_TEXT)
        sys.exit(0)

    cmd = sys.argv[1]

    try:
        cfg = load_config()
    except (FileNotFoundError, json.JSONDecodeError) as exc:
        print(json.dumps({"ok": False, "error": f"Config error: {exc}"}, indent=2))
        sys.exit(1)

    if cmd == "health":
        result = cmd_health(cfg)
    elif cmd == "create-collection":
        result = cmd_create_collection(cfg)
    elif cmd == "search":
        if len(sys.argv) < 3:
            result = {"ok": False, "error": "search requires a query text argument"}
        else:
            result = cmd_search(cfg, sys.argv[2])
    elif cmd == "index":
        if len(sys.argv) < 4:
            result = {"ok": False, "error": "index requires <id> and <text> arguments"}
        else:
            metadata = parse_metadata(sys.argv[4:])
            result = cmd_index(cfg, sys.argv[2], sys.argv[3], metadata)
    elif cmd == "index-file":
        if len(sys.argv) < 3:
            result = {"ok": False, "error": "index-file requires a <filepath> argument"}
        else:
            result = cmd_index_file(cfg, sys.argv[2])
    elif cmd == "count":
        result = cmd_count(cfg)
    else:
        result = {"ok": False, "error": f"Unknown command: {cmd}. Use --help for usage."}

    print(json.dumps(result, indent=2))
    sys.exit(0 if result.get("ok") else 1)


if __name__ == "__main__":
    main()