"""Build the vector store from rag/sources.yaml.

    python -m rag.ingest                 # all sources
    python -m rag.ingest --only mujoco   # one collection
    python -m rag.ingest --reset         # drop and rebuild collections

Clones each source repository (shallow, sparse) into <data>/sources/, chunks
the included files, and upserts them into Chroma at <data>/chroma/. Also
writes <data>/ros2_interfaces.json, the catalogue of message/service types
per distro that /code/generate validates against.
"""

from __future__ import annotations

import argparse
import fnmatch
import json
import subprocess
import sys
from pathlib import Path

import yaml

from api.config import get_settings

from .chunking import Chunk, chunk_file, interface_name, parse_mjcf_schema
from .store import COLLECTIONS, RagStore

SOURCES_FILE = Path(__file__).with_name("sources.yaml")


def load_sources(path: Path = SOURCES_FILE) -> list[dict]:
    return yaml.safe_load(path.read_text())["sources"]


def _git(*args: str, cwd: Path | None = None) -> str:
    return subprocess.run(["git", *args], cwd=cwd, check=True, capture_output=True, text=True).stdout


def checkout(source: dict, sources_dir: Path) -> Path:
    """Shallow sparse clone (or update) of a source repo; returns its directory."""
    dest = sources_dir / source["name"]
    ref = source.get("ref", "main")
    if not (dest / ".git").exists():
        dest.parent.mkdir(parents=True, exist_ok=True)
        _git(
            "clone",
            "--quiet",
            "--depth",
            "1",
            "--filter=blob:none",
            "--no-checkout",
            "-b",
            ref,
            source["repo"],
            str(dest),
        )
    else:
        _git("fetch", "--quiet", "--depth", "1", "origin", ref, cwd=dest)
        _git("reset", "--quiet", "--hard", "FETCH_HEAD", cwd=dest)
    patterns = list(source.get("checkout", [])) + list(source["include"])
    _git("sparse-checkout", "set", "--no-cone", *patterns, cwd=dest)
    _git("checkout", "--quiet", cwd=dest)
    return dest


def included_files(root: Path, patterns: list[str]) -> list[str]:
    files = []
    for p in sorted(root.rglob("*")):
        if p.is_file() and ".git" not in p.parts:
            rel = p.relative_to(root).as_posix()
            if any(fnmatch.fnmatch(rel, pat) for pat in patterns):
                files.append(rel)
    return files


def ingest(only: str | None = None, reset: bool = False, log=print) -> dict:
    settings = get_settings()
    sources_dir = settings.data_dir / "sources"
    store = RagStore(settings.data_dir / "chroma")
    sources = [s for s in load_sources() if only in (None, s["collection"])]
    if reset:
        for name in {s["collection"] for s in sources}:
            store.reset(name)

    interfaces: dict[str, dict[str, str]] = {}
    catalog_path = settings.data_dir / "ros2_interfaces.json"
    if catalog_path.exists() and not reset:
        interfaces = json.loads(catalog_path.read_text())

    counts: dict[str, int] = {}
    for source in sources:
        log(f"[{source['name']}] fetching {source['repo']}@{source.get('ref', 'main')}")
        root = checkout(source, sources_dir)
        files = included_files(root, source["include"])
        if not files:
            log(f"[{source['name']}] warning: no files matched {source['include']}")
        chunks: list[Chunk] = []
        meta = dict(source.get("meta", {}))
        for rel in files:
            text = (root / rel).read_text(errors="replace")
            chunks += chunk_file(rel, text, source["name"], meta, source.get("package"))
            if rel.endswith("XMLschema.rst"):
                schema_path = settings.data_dir / "mjcf_schema.json"
                schema_path.write_text(json.dumps(parse_mjcf_schema(text), indent=1))
                log(f"[{source['name']}] wrote {schema_path}")
            name = interface_name(rel, source.get("package"))
            if name and meta.get("distro"):
                interfaces.setdefault(meta["distro"], {})[name] = text
        n = store.add(source["collection"], chunks)
        counts[source["name"]] = n
        log(f"[{source['name']}] {len(files)} files -> {n} chunks")

    if interfaces:
        catalog_path.parent.mkdir(parents=True, exist_ok=True)
        catalog_path.write_text(json.dumps(interfaces, indent=1, sort_keys=True))
        log(f"interface catalogue: {', '.join(f'{d}={len(v)}' for d, v in sorted(interfaces.items()))}")
    for name in COLLECTIONS:
        log(f"collection {name}: {store.count(name)} chunks")
    return counts


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--only", choices=COLLECTIONS, help="ingest one collection")
    parser.add_argument("--reset", action="store_true", help="drop the collections first")
    args = parser.parse_args(argv)
    try:
        ingest(args.only, args.reset)
    except subprocess.CalledProcessError as e:
        print(f"error: {' '.join(e.cmd)} failed:\n{e.stderr}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())