"""Edge Impulse docs RAG — retrieval + grounded generation.

Retrieval: FAISS (inner-product) over the prebuilt index in ``data/index`` using
the same ``all-MiniLM-L6-v2`` sentence embedder the index was built with.

Generation: the published quantized model ``edgeimpulse/edgeimpulse-docs-qwen-0.5b``
served through any OpenAI-compatible endpoint — e.g. llama.cpp's ``llama-server``
or Ollama. Only the tiny GGUF is needed for generation, so no training stack is
required to run this assistant.

The raw document corpus and the index-building pipeline are intentionally not
part of this repository; the prebuilt index is all you need at inference time.
"""
from __future__ import annotations

import argparse
import json
import os
import pickle
from functools import lru_cache
from pathlib import Path
from typing import Any

import faiss
import requests
from sentence_transformers import SentenceTransformer

DEFAULT_INDEX_DIR = Path(os.environ.get("RAG_INDEX_DIR", "data/index"))

# OpenAI-compatible generation endpoint (llama.cpp `llama-server` or Ollama).
#   llama.cpp : llama-server -m qwen-edgeai-q4_k_m.gguf --port 8080 --jinja
#   ollama    : ollama run hf.co/edgeimpulse/edgeimpulse-docs-qwen-0.5b
DEFAULT_API_BASE = os.environ.get("RAG_API_BASE", "http://127.0.0.1:8080/v1")
DEFAULT_MODEL = os.environ.get("RAG_MODEL", "edgeimpulse/edgeimpulse-docs-qwen-0.5b")
DEFAULT_API_KEY = os.environ.get("RAG_API_KEY", "sk-no-key-required")

SYSTEM_PROMPT = (
    "You are an Edge Impulse documentation assistant. Answer only from the "
    "provided context. If the context does not contain the answer, say what is "
    "missing and suggest the closest relevant docs source. Be concise."
)


@lru_cache(maxsize=1)
def load_retriever(index_dir: str):
    root = Path(index_dir)
    metadata = json.loads((root / "metadata.json").read_text(encoding="utf-8"))
    index = faiss.read_index(str(root / "edge_impulse_docs.faiss"))
    with (root / "chunks.pkl").open("rb") as f:
        chunks = pickle.load(f)
    embedder = SentenceTransformer(metadata["embedding_model"])
    return index, chunks, embedder, metadata


def retrieve(question: str, index_dir: Path = DEFAULT_INDEX_DIR, k: int = 4) -> list[dict[str, Any]]:
    index, chunks, embedder, _ = load_retriever(str(index_dir))
    q_emb = embedder.encode(
        [question], convert_to_numpy=True, normalize_embeddings=True
    ).astype("float32")
    scores, ids = index.search(q_emb, k)
    results: list[dict[str, Any]] = []
    for score, idx in zip(scores[0], ids[0]):
        if idx < 0:
            continue
        record = dict(chunks[int(idx)])
        record["score"] = float(score)
        results.append(record)
    return results


def build_messages(question: str, contexts: list[dict[str, Any]]) -> list[dict[str, str]]:
    context_text = "\n\n".join(
        f"Source: {item['source']}\n{item['text']}" for item in contexts
    )
    user = f"Context:\n{context_text}\n\nQuestion: {question}"
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user},
    ]


def generate(
    messages: list[dict[str, str]],
    api_base: str = DEFAULT_API_BASE,
    model: str = DEFAULT_MODEL,
    api_key: str = DEFAULT_API_KEY,
    max_new_tokens: int = 320,
) -> str:
    payload = {
        "model": model,
        "messages": messages,
        "temperature": 0.3,
        "top_p": 0.9,
        "max_tokens": max_new_tokens,
        # Honoured by llama.cpp's server; ignored by backends that don't support it.
        "repeat_penalty": 1.2,
    }
    resp = requests.post(
        f"{api_base.rstrip('/')}/chat/completions",
        headers={"Authorization": f"Bearer {api_key}"},
        json=payload,
        timeout=120,
    )
    resp.raise_for_status()
    return resp.json()["choices"][0]["message"]["content"].strip()


def ask(
    question: str,
    index_dir: Path = DEFAULT_INDEX_DIR,
    k: int = 4,
    max_new_tokens: int = 320,
    no_generate: bool = False,
    api_base: str = DEFAULT_API_BASE,
    model: str = DEFAULT_MODEL,
) -> str:
    contexts = retrieve(question, index_dir, k)
    sources = "\n".join(f"- {item['source']} ({item['score']:.3f})" for item in contexts)
    if no_generate:
        return "Retrieved context:\n" + sources
    answer = generate(
        build_messages(question, contexts),
        api_base=api_base,
        model=model,
        max_new_tokens=max_new_tokens,
    )
    return f"{answer}\n\nSources:\n{sources}"


def main() -> None:
    parser = argparse.ArgumentParser(description="Ask the Edge Impulse docs RAG assistant.")
    parser.add_argument("question")
    parser.add_argument("--index-dir", type=Path, default=DEFAULT_INDEX_DIR)
    parser.add_argument("--k", type=int, default=4)
    parser.add_argument("--max-new-tokens", type=int, default=320)
    parser.add_argument("--api-base", default=DEFAULT_API_BASE)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--no-generate", action="store_true", help="Only print retrieved chunks.")
    args = parser.parse_args()
    print(
        ask(
            args.question,
            index_dir=args.index_dir,
            k=args.k,
            max_new_tokens=args.max_new_tokens,
            no_generate=args.no_generate,
            api_base=args.api_base,
            model=args.model,
        )
    )


if __name__ == "__main__":
    main()