"""Local-folder RAG ingest (Mode B) — runs on the USER'S machine via the
@melaya/runner subprocess. Reads docs from a local folder, embeds via the
user's local Ollama (or LM Studio), persists the Qdrant vector store at
`~/.melaya-runner/rag/<pipeline_name>/store/`.

Data residency guarantee:
  * Document text is NEVER uploaded to the prod box
  * Embedding API calls go to localhost (user's Ollama at :11434)
  * Vector store sits at ~/.melaya-runner/... on the user's filesystem
  * The only network egress is the Socket.IO progress events
    (file names + chunk counts only, NO chunk text)

Invoked by packages/runner/src/connection.ts on `rag:ingest` Socket.IO event.
Emits `progress: <json>` lines to stdout that the parent process forwards to
the FE so the user sees per-file ingest progress in real time.
"""
from __future__ import annotations

import argparse
import asyncio
import datetime as _dt
import hashlib
import json
import os
import sys
import time
from pathlib import Path

# Path hack — the runner's spawned venv has agentscope vendored at the
# PYTHONPATH it sets (the parent connection.ts sets env.PYTHONPATH to the
# shared bundle dir). Confirm imports work, fail loud otherwise.
try:
    from agentscope.rag import (  # noqa: F401
        TextReader, PDFReader, WordReader, PowerPointReader, ExcelReader,
    )
except ImportError as exc:
    print(f"ImportError: agentscope.rag not reachable on PYTHONPATH={os.environ.get('PYTHONPATH')!r}: {exc}", file=sys.stderr)
    sys.exit(2)

# Ensure NLTK's sentence tokenizer data is on disk. agentscope's readers
# do `nltk.download(..., quiet=True)` inside the chunking call but with
# quiet=True a network failure surfaces only as "Resource 'punkt_tab' not
# found" at sent_tokenize time. We force the download once up-front,
# loudly, so the user sees a clear network-error if it fails — beats
# 240 per-file failures with the same opaque message.
def _ensure_nltk_data() -> None:
    try:
        import nltk  # type: ignore
    except ImportError:
        print("nltk missing — runner venv needs `pip install nltk`", file=sys.stderr)
        return
    needed = ["punkt", "punkt_tab"]
    for pkg in needed:
        try:
            nltk.data.find(f"tokenizers/{pkg}")
        except LookupError:
            try:
                ok = nltk.download(pkg, quiet=True)
                if not ok:
                    print(f"nltk.download({pkg!r}) returned False — likely offline. "
                          f"Run `python -m nltk.downloader {pkg}` once with network.",
                          file=sys.stderr)
            except Exception as exc:  # noqa: BLE001
                print(f"nltk.download({pkg!r}) failed: {type(exc).__name__}: {exc}",
                      file=sys.stderr)

_ensure_nltk_data()


def _emit(kind: str, **fields):
    """Stream a progress event to the runner subprocess parent which
    forwards it to the FE over Socket.IO."""
    payload = {"kind": kind, **fields, "ts": _dt.datetime.utcnow().isoformat() + "Z"}
    print("progress: " + json.dumps(payload), flush=True)


SUPPORTED_EXTS = {".txt", ".md", ".csv", ".json", ".pdf",
                  ".docx", ".pptx", ".xlsx"}


def _reader_for(ext: str):
    """Map extension to agentscope reader. Returns None when no extractor
    exists (legacy .doc, image-only PDFs that need OCR, etc.)."""
    return {
        ".txt":  TextReader, ".md":   TextReader,
        ".csv":  TextReader, ".json": TextReader,
        ".pdf":  PDFReader,
        ".docx": WordReader,
        ".pptx": PowerPointReader,
        ".xlsx": ExcelReader,
    }.get(ext)


def _resolve_store_path(pipeline_name: str) -> Path:
    """Mode B store path — under ~/.melaya-runner/rag/<pipeline>/store/.
    Matches the server-side facade's resolve_store_path(mode="local_folder")
    so cross-runtime debugging is consistent."""
    return Path.home() / ".melaya-runner" / "rag" / pipeline_name / "store"


def _resolve_manifest_path(pipeline_name: str) -> Path:
    return _resolve_store_path(pipeline_name).parent / ".ingest.json"


async def _ingest(args) -> int:
    folder = Path(args.folder).expanduser()
    if not folder.exists() or not folder.is_dir():
        _emit("error", reason=f"folder not found or not a directory: {folder}")
        return 3

    # Build the facade-equivalent: pick the embedder, instantiate Qdrant.
    # We import lazily so a missing dep produces a clear error message.
    try:
        if args.embedder_provider == "ollama":
            from agentscope.embedding import OllamaTextEmbedding
            # Centralized dimension lookup — same as the server-side facade
            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(args.embedder_model.lower())
            if dim is None:
                _emit("error", reason=(
                    f"unknown ollama embedding model '{args.embedder_model}' — "
                    f"supported: {sorted(DIMS)}"
                ))
                return 4
            embedder = OllamaTextEmbedding(
                model_name=args.embedder_model,
                dimensions=dim,
                host=None,  # → uses default localhost:11434
            )
        elif args.embedder_provider == "lmstudio":
            from agentscope.embedding import OpenAITextEmbedding
            # LM Studio exposes an OpenAI-compatible /v1; reuse the OpenAI
            # client pointed at its local endpoint. Dim default = 768 (the
            # FE picker should report the real dim per loaded model).
            # base_url MUST be passed at construction (the OpenAI 1.x
            # AsyncClient bakes the URL into its httpx client during
            # __init__; setting `client.base_url = ...` after the fact is a
            # no-op and the next request resolves an empty hostname →
            # `[Errno 8] nodename nor servname provided`).
            embedder = OpenAITextEmbedding(
                api_key="lmstudio",
                model_name=args.embedder_model,
                dimensions=768,
                base_url="http://localhost:1234/v1",
            )
            # Diagnostic: surface the URLs the httpx client will actually
            # hit. Prior runs were failing with `[Errno 8] nodename nor
            # servname provided` — the high-level `.base_url` attribute
            # was being set but the underlying httpx client kept the
            # default `api.openai.com` from construction. With the kwarg
            # path these should both be `localhost:1234/v1`.
            try:
                _cb = getattr(embedder.client, "base_url", "?")
                _hb = getattr(getattr(embedder.client, "_client", None), "base_url", "?")
                _emit("lmstudio_init", client_base_url=str(_cb), httpx_base_url=str(_hb))
            except Exception as _exc:
                _emit("lmstudio_init", error=f"{type(_exc).__name__}: {_exc}")
            # Probe the LM Studio HTTP server BEFORE the first embed call
            # so a misconfigured/off server surfaces as one clear error
            # rather than 240 cryptic per-file failures.
            try:
                import urllib.request, json as _json
                with urllib.request.urlopen(
                    "http://localhost:1234/v1/models", timeout=3
                ) as _r:
                    _models = _json.loads(_r.read().decode("utf-8"))
                    _names = [m.get("id") for m in _models.get("data", [])]
                    _emit("lmstudio_probe", ok=True, models=_names)
                    if args.embedder_model not in _names:
                        _emit("warning", reason=(
                            f"LM Studio is reachable but model "
                            f"'{args.embedder_model}' is not loaded. "
                            f"Loaded: {_names}"
                        ))
            except Exception as _exc:
                _emit("lmstudio_probe", ok=False,
                      error=f"{type(_exc).__name__}: {_exc}")
                _emit("error", reason=(
                    f"Cannot reach LM Studio at http://localhost:1234. "
                    f"Make sure the LM Studio app's local server is started "
                    f"(Developer → Start Server) and the model "
                    f"'{args.embedder_model}' is loaded. Underlying error: "
                    f"{type(_exc).__name__}: {_exc}"
                ))
                return 6
            # LM Studio's /v1/embeddings rejects the OpenAI-specific
            # `dimensions` kwarg (only OpenAI text-embedding-3 supports
            # output-dim reduction). Strip it before each call so the
            # request is shaped exactly as LMS expects.
            _orig_create = embedder.client.embeddings.create
            async def _lms_create(**kw):
                kw.pop("dimensions", None)
                return await _orig_create(**kw)
            embedder.client.embeddings.create = _lms_create  # type: ignore[assignment]
        else:
            _emit("error", reason=(
                f"refused embedder provider '{args.embedder_provider}' for Mode B "
                "— only ollama and lmstudio keep documents private"
            ))
            return 5

        from agentscope.rag import QdrantStore, SimpleKnowledge
        store_path = _resolve_store_path(args.pipeline_name)
        store_path.mkdir(parents=True, exist_ok=True)
        # ── Fix qdrant-client's location-vs-path ambiguity ───────────────
        # agentscope's QdrantStore only exposes `location=` and forwards
        # it verbatim to `AsyncQdrantClient(location=...)`. When `location`
        # is a filesystem path (not `:memory:` and not a URL), qdrant-client
        # tries to parse it as a host:port → getaddrinfo → fails with
        # `[Errno 8] nodename nor servname provided` on Mac / `[Errno
        # 11001] getaddrinfo failed` on Windows. Per-file embedding then
        # crashes 240 times in a row. Reroute `location=<path>` into the
        # dedicated `path=` kwarg the embedded-mode wants.
        import qdrant_client as _qc
        _orig_qcinit = _qc.AsyncQdrantClient.__init__
        def _patched_qcinit(self, location=None, **_kw):
            if (location and location != ":memory:"
                    and "://" not in str(location)
                    and not _kw.get("path")):
                _kw["path"] = location
                location = None
            _orig_qcinit(self, location=location, **_kw)
        _qc.AsyncQdrantClient.__init__ = _patched_qcinit
        store = QdrantStore(
            location=str(store_path),
            collection_name=args.pipeline_name,
            dimensions=embedder.dimensions,
            distance="Cosine",
        )
        kb = SimpleKnowledge(embedding_store=store, embedding_model=embedder)
    except Exception as exc:
        _emit("error", reason=f"failed to build KB: {type(exc).__name__}: {exc}")
        return 6

    # Read the manifest to skip unchanged files
    manifest_path = _resolve_manifest_path(args.pipeline_name)
    if manifest_path.exists():
        try:
            manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        except Exception:
            manifest = {"files": {}, "embedder": None}
    else:
        manifest = {"files": {}, "embedder": None}
    stored = manifest.get("embedder") or {}
    embedder_changed = (
        stored.get("provider") != args.embedder_provider
        or stored.get("model") != args.embedder_model
    )
    if embedder_changed:
        _emit("embedder_changed", was=stored, now={
            "provider": args.embedder_provider, "model": args.embedder_model,
        })

    started = time.time()
    chunks_added = 0
    files_processed = 0
    skipped = 0
    files_failed = []

    # Walk the folder — only top-level + one level deep. Going arbitrarily
    # deep would silently include things like venv/.git which would be
    # surprising. If someone needs recursive, they can manually flatten.
    candidates = [p for p in sorted(folder.iterdir()) if p.is_file()]
    candidates += [p for sub in folder.iterdir() if sub.is_dir() and not sub.name.startswith(".")
                   for p in sorted(sub.iterdir()) if p.is_file()]

    total = sum(1 for p in candidates if p.suffix.lower() in SUPPORTED_EXTS)
    _emit("scanned", folder=str(folder), total_supported=total)

    for f in candidates:
        ext = f.suffix.lower()
        if ext not in SUPPORTED_EXTS:
            continue
        # macOS AppleDouble resource forks (._foo.docx, ._.DS_Store, …) —
        # they're not real documents and every extractor will choke on them.
        if f.name.startswith("._") or f.name == ".DS_Store":
            continue
        try:
            raw = f.read_bytes()
        except Exception as exc:
            files_failed.append({"file": str(f), "error": f"read: {type(exc).__name__}: {exc}"})
            continue
        sha = hashlib.sha256(raw).hexdigest()
        prev = manifest["files"].get(f.name, {})
        if (not embedder_changed) and prev.get("sha256") == sha:
            skipped += 1
            _emit("skipped_sha_clean", file=f.name)
            continue
        reader_cls = _reader_for(ext)
        if reader_cls is None:
            files_failed.append({"file": str(f), "error": "no extractor"})
            continue
        try:
            reader = reader_cls()
            docs = await reader(str(f))
            for doc in docs:
                if hasattr(doc.metadata, "doc_id"):
                    doc.metadata.doc_id = f.name
            await kb.add_documents(docs)
            chunks_added += len(docs)
            manifest["files"][f.name] = {
                "sha256": sha, "chunks": len(docs),
                "ingested_at": _dt.datetime.utcnow().isoformat() + "Z",
            }
            files_processed += 1
            _emit("file_done", file=f.name, chunks=len(docs))
        except Exception as exc:
            files_failed.append({
                "file": str(f), "error": f"{type(exc).__name__}: {exc}",
            })
            _emit("file_failed", file=f.name, error=str(exc))

    manifest["embedder"] = {
        "provider": args.embedder_provider, "model": args.embedder_model,
    }
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")

    elapsed_ms = int((time.time() - started) * 1000)
    _emit("done",
          files_processed=files_processed,
          chunks_added=chunks_added,
          skipped=skipped,
          embedder_changed=embedder_changed,
          elapsed_ms=elapsed_ms,
          errors=files_failed)
    return 0


def _build_embedder(provider: str, model: str):
    """Embedder factory shared between Mode B (local folder) and Mode C
    (cloud files via HTTP). Returns the agentscope embedder instance +
    a probe error message string when something is misconfigured (empty
    string means OK)."""
    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:
            return None, (
                f"unknown ollama embedding model '{model}' — supported: {sorted(DIMS)}"
            )
        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",
        )
        # Strip OpenAI-specific `dimensions` kwarg that LMS rejects
        _orig_create = embedder.client.embeddings.create
        async def _lms_create(**kw):
            kw.pop("dimensions", None)
            return await _orig_create(**kw)
        embedder.client.embeddings.create = _lms_create  # type: ignore[assignment]
        return embedder, ""
    return None, f"unsupported embedder provider '{provider}'"


async def _ingest_cloud_files(args) -> int:
    """Mode C ingest.

    Files originate on prod but the runner does the embedding so the
    operator can use a free local model (ollama / lmstudio) without
    needing an OpenAI key. Both the file CONTENTS and the resulting
    (chunks, vectors) tuples flow over Socket.IO end-to-end:

      • The server (tRPC bridge in runnerNamespace.ts) reads the files
        from prod's <pipeline_dir>/rag/retrieval/, base64-encodes them,
        and bakes them into the `rag:ingest` event payload. The runner
        connection.ts unpacks them into ``<tmp>/files-dir/`` and
        passes ``--files-dir <tmp> --manifest-json <tmp>/manifest.json``
        here.
      • For each file we parse + chunk + embed locally, then emit a
        ``progress: {"kind":"upsert_chunks", ...}`` line to stdout. The
        runner forwards that line as a ``rag:ingest-progress`` Socket.IO
        event; the server's progress handler recognises the
        ``upsert_chunks`` kind and POSTs the payload to
        ``http://127.0.0.1:8766/pipelines/{name}/docs/retrieval/upsert``
        — the loopback-only builder API that owns the QdrantStore.
      • Nothing on the wire ever needs to leave the user's machine
        via HTTP. The whole transport is the already-authenticated
        Socket.IO channel, which means we don't have to expose any
        builder endpoints publicly.
    """
    files_dir = Path(args.files_dir)
    if not files_dir.exists() or not files_dir.is_dir():
        _emit("error", reason=f"files-dir not found: {files_dir}")
        return 3

    embedder, err = _build_embedder(args.embedder_provider, args.embedder_model)
    if not embedder:
        _emit("error", reason=err)
        return 4

    manifest: dict = {}
    if args.manifest_json:
        try:
            manifest = json.loads(Path(args.manifest_json).read_text(encoding="utf-8"))
        except Exception as exc:
            _emit("warning", reason=f"manifest read failed: {exc}")
    files: list[dict] = manifest.get("files") or []
    stored_embedder = manifest.get("stored_embedder") or {}
    previously_ingested = manifest.get("previously_ingested") or {}
    embedder_changed = (
        stored_embedder.get("provider") != args.embedder_provider
        or stored_embedder.get("model") != args.embedder_model
    )
    _emit("started", mode="cloud_files", pipeline=args.pipeline_name,
          embedder_changed=embedder_changed, total_supported=len(files))

    started = time.time()
    chunks_added_total = 0
    files_processed = 0
    skipped = 0
    files_failed: list[dict] = []

    embedder_payload = {
        "provider": args.embedder_provider,
        "model": args.embedder_model,
        "dimensions": embedder.dimensions,
    }

    first_emit = True
    for entry in files:
        name = entry["name"]
        sha = entry["sha256"]
        ext = (entry.get("ext") or "").lower()
        if ext not in SUPPORTED_EXTS:
            files_failed.append({"file": name, "error": f"unsupported ext '{ext}'"})
            _emit("file_failed", file=name, error=f"unsupported ext '{ext}'")
            continue
        if not embedder_changed and previously_ingested.get(name, {}).get("sha256") == sha:
            skipped += 1
            _emit("skipped_sha_clean", file=name)
            continue
        tmp_path = files_dir / name
        if not tmp_path.is_file():
            files_failed.append({"file": name, "error": "file missing from cloud-files dir"})
            _emit("file_failed", file=name, error="missing locally")
            continue
        reader_cls = _reader_for(ext)
        if reader_cls is None:
            files_failed.append({"file": name, "error": "no extractor"})
            _emit("file_failed", file=name, error="no extractor")
            continue
        try:
            reader = reader_cls()
            docs = await reader(str(tmp_path))
            for i, doc in enumerate(docs):
                if hasattr(doc.metadata, "doc_id"):
                    doc.metadata.doc_id = name
                if hasattr(doc.metadata, "chunk_id"):
                    doc.metadata.chunk_id = i
            texts: list[str] = []
            for doc in docs:
                content = getattr(doc.metadata, "content", None)
                text = (
                    (content or {}).get("text", "")
                    if isinstance(content, dict)
                    else getattr(content, "text", "")
                )
                if text:
                    texts.append(text)
            if not texts:
                files_processed += 1
                _emit("file_done", file=name, chunks=0)
                _emit("upsert_chunks",
                      embedder=embedder_payload,
                      chunks=[],
                      files_done=[{"name": name, "sha256": sha, "chunks": 0}],
                      embedder_changed=embedder_changed and first_emit,
                      finalize=False)
                first_emit = False
                continue
            vectors_resp = await embedder(text=texts)
            # agentscope's embedding adapters already unwrap to a flat list
            # of vectors (list[list[float]]) via _openai_embedding.py L104 /
            # _ollama_embedding.py — do NOT re-access `.embedding`.
            vectors = list(vectors_resp.embeddings)
            if len(vectors) != len(texts):
                raise RuntimeError(
                    f"embedder returned {len(vectors)} vectors for {len(texts)} chunks"
                )
            upsert_chunks = [
                {
                    "text": text,
                    "vector": [float(v) for v in vec],
                    "metadata": {
                        "doc_id": name,
                        "chunk_id": i,
                        "total_chunks": len(texts),
                    },
                }
                for i, (text, vec) in enumerate(zip(texts, vectors))
            ]
            # Emit one upsert event per file. Server forwards to builder
            # /upsert. For small per-file payloads this is well below the
            # Socket.IO default 1 MB/event limit (60-doc corpora ≈ 6 MB
            # spread across ~60 events).
            _emit("upsert_chunks",
                  embedder=embedder_payload,
                  chunks=upsert_chunks,
                  files_done=[{"name": name, "sha256": sha, "chunks": len(texts)}],
                  embedder_changed=embedder_changed and first_emit,
                  finalize=False)
            first_emit = False
            chunks_added_total += len(texts)
            files_processed += 1
            _emit("file_done", file=name, chunks=len(texts))
        except Exception as exc:
            files_failed.append({
                "file": name, "error": f"{type(exc).__name__}: {exc}",
            })
            _emit("file_failed", file=name, error=str(exc))

    # Finalize: tells the server to stamp the embedder identity into the
    # manifest. Sent as one final, empty-chunk upsert event.
    _emit("upsert_chunks",
          embedder=embedder_payload,
          chunks=[],
          files_done=[],
          embedder_changed=False,
          finalize=True)

    elapsed_ms = int((time.time() - started) * 1000)
    _emit("done",
          mode="cloud_files",
          files_processed=files_processed,
          chunks_added=chunks_added_total,
          skipped=skipped,
          embedder_changed=embedder_changed,
          elapsed_ms=elapsed_ms,
          errors=files_failed)
    return 0


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", choices=["local_folder", "cloud_files"],
                    default="local_folder")
    ap.add_argument("--pipeline-name", required=True)
    ap.add_argument("--folder", default="")
    ap.add_argument("--embedder-provider", required=True, choices=["ollama", "lmstudio"])
    ap.add_argument("--embedder-model", required=True)
    # Mode C only (cloud_files): files were base64-encoded by the
    # server into the rag:ingest payload and the runner's connection.ts
    # unpacked them to disk at --files-dir. --manifest-json points at
    # a small JSON manifest (file list + sha + previously-ingested map
    # so we can skip unchanged files).
    ap.add_argument("--files-dir", default="",
                    help="Mode C only — directory containing the unpacked file payload.")
    ap.add_argument("--manifest-json", default="",
                    help="Mode C only — path to manifest JSON (file list, prior ingest, embedder).")
    args = ap.parse_args()
    if args.mode == "cloud_files":
        if not args.files_dir or not args.manifest_json:
            _emit("error", reason="cloud_files mode requires --files-dir + --manifest-json")
            sys.exit(2)
        sys.exit(asyncio.run(_ingest_cloud_files(args)))
    else:
        if not args.folder:
            _emit("error", reason="local_folder mode requires --folder")
            sys.exit(2)
        sys.exit(asyncio.run(_ingest(args)))


if __name__ == "__main__":
    main()