"""Runner-side RAG retrieval test (local embedder paths).

Two modes, selected by --embed-only:
  * Mode B (local folder, default): the vector store lives on THIS machine
    (~/.melaya-runner/rag/<pipeline>/store), so we embed the query AND search
    locally, returning the matched chunks.
  * Mode C (cloud files, --embed-only): the chunks were upserted to the
    PROD-side store during ingest, so they are NOT on this box. We only embed
    the query with the user's local embedder and return the raw vector; the
    server searches the prod store by that vector (see builder
    /docs/retrieval/search_by_vector + runnerNamespace forwarding).

Either way the embedder is local (ollama / lmstudio), which is why this can't
run server-side: prod has no LM Studio / Ollama to embed the query.
"""
from __future__ import annotations

import argparse
import asyncio
import json
import sys
from pathlib import Path


def _resolve_store_path(pipeline_name: str) -> Path:
    return Path.home() / ".melaya-runner" / "rag" / pipeline_name / "store"


def _build_embedder(provider: str, model: str):
    if provider == "ollama":
        from agentscope.embedding import OllamaTextEmbedding
        dims = {
            "nomic-embed-text": 768,
            "nomic-embed-text:v1.5": 768,
            "mxbai-embed-large": 1024,
            "all-minilm": 384,
            "snowflake-arctic-embed": 1024,
        }
        dim = dims.get(model.lower())
        if dim is None:
            raise ValueError(f"unknown ollama embedding model '{model}'")
        return OllamaTextEmbedding(model_name=model, dimensions=dim, host=None)
    if provider == "lmstudio":
        from agentscope.embedding import OpenAITextEmbedding
        embedder = OpenAITextEmbedding(
            api_key="lmstudio",
            model_name=model,
            dimensions=768,
            base_url="http://localhost:1234/v1",
        )
        original_create = embedder.client.embeddings.create

        async def _lms_create(**kwargs):
            kwargs.pop("dimensions", None)
            return await original_create(**kwargs)

        embedder.client.embeddings.create = _lms_create  # type: ignore[assignment]
        return embedder
    raise ValueError(f"unsupported embedder provider '{provider}'")


def _extract_text(doc) -> str:
    content = getattr(doc.metadata, "content", None)
    if isinstance(content, dict):
        return str(content.get("text") or "")
    return str(getattr(content, "text", "") or "")


async def _embed_only(args) -> dict:
    """Mode C: embed the query with the user's LOCAL embedder and return the
    raw vector. The vectors live in the prod-side store (Mode C ingest upserted
    them there), so the SERVER does the search — this box only embeds. No store
    is opened here."""
    from agentscope.message import TextBlock
    embedder = _build_embedder(args.embedder_provider, args.embedder_model)
    res = await embedder([TextBlock(type="text", text=str(args.query or "")[:2000])])
    vector = list(res.embeddings[0])
    return {"ok": True, "vector": vector}


async def _retrieve(args) -> dict:
    if getattr(args, "embed_only", False):
        return await _embed_only(args)
    store_path = _resolve_store_path(args.pipeline_name)
    if not store_path.exists():
        return {
            "ok": False,
            "error": f"runner RAG store not found: {store_path}",
            "results": [],
        }

    import qdrant_client as _qc
    if not getattr(_qc.AsyncQdrantClient.__init__, "_melaya_patched", False):
        original_init = _qc.AsyncQdrantClient.__init__

        def _patched_init(self, location=None, **kwargs):
            if (
                location
                and location != ":memory:"
                and "://" not in str(location)
                and not kwargs.get("path")
            ):
                kwargs["path"] = location
                location = None
            original_init(self, location=location, **kwargs)

        _patched_init._melaya_patched = True  # type: ignore[attr-defined]
        _qc.AsyncQdrantClient.__init__ = _patched_init  # type: ignore[assignment]

    from agentscope.rag import QdrantStore, SimpleKnowledge

    embedder = _build_embedder(args.embedder_provider, args.embedder_model)
    store = QdrantStore(
        location=str(store_path),
        collection_name=args.pipeline_name,
        dimensions=embedder.dimensions,
        distance="Cosine",
    )
    kb = SimpleKnowledge(embedding_store=store, embedding_model=embedder)
    docs = await kb.retrieve(
        query=str(args.query or "")[:2000],
        limit=max(1, min(20, int(args.limit))),
    )
    results = []
    for doc in docs:
        score = getattr(doc, "score", None)
        results.append({
            "source": getattr(doc.metadata, "doc_id", "?"),
            "chunk_id": getattr(doc.metadata, "chunk_id", "?"),
            "score": score,
            "text": _extract_text(doc),
        })
    return {
        "ok": True,
        "store_path": str(store_path),
        "results": results,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--pipeline-name", required=True)
    parser.add_argument("--embedder-provider", required=True, choices=["ollama", "lmstudio"])
    parser.add_argument("--embedder-model", required=True)
    parser.add_argument("--query", required=True)
    parser.add_argument("--limit", type=int, default=5)
    # Mode C: embed the query and print {ok, vector}; do not open a local store.
    parser.add_argument("--embed-only", dest="embed_only", action="store_true")
    args = parser.parse_args()
    try:
        print(json.dumps(asyncio.run(_retrieve(args)), ensure_ascii=False))
    except Exception as exc:  # noqa: BLE001
        print(json.dumps({
            "ok": False,
            "error": f"{type(exc).__name__}: {exc}",
            "results": [],
        }), file=sys.stdout)
        sys.exit(1)


if __name__ == "__main__":
    main()