from __future__ import annotations

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

import faiss
from peft import PeftModel
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline


DEFAULT_INDEX_DIR = Path("data/index")
DEFAULT_BASE_MODEL = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
DEFAULT_ADAPTER = ""


@lru_cache(maxsize=1)
def load_retriever(index_dir: str) -> tuple[Any, list[dict[str, Any]], SentenceTransformer, dict[str, Any]]:
    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


@lru_cache(maxsize=1)
def load_generator(base_model: str, adapter: str):
    tokenizer = AutoTokenizer.from_pretrained(base_model)
    model = AutoModelForCausalLM.from_pretrained(base_model, device_map="auto")
    if adapter:
        model = PeftModel.from_pretrained(model, adapter)
    return pipeline("text-generation", model=model, tokenizer=tokenizer)


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_prompt(question: str, contexts: list[dict[str, Any]]) -> str:
    context_text = "\n\n".join(
        f"Source: {item['source']}\n{item['text']}" for item in contexts
    )
    return (
        "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.\n\n"
        f"Context:\n{context_text}\n\n"
        f"Question: {question}\n"
        "Answer:"
    )


def ask(
    question: str,
    index_dir: Path = DEFAULT_INDEX_DIR,
    base_model: str = DEFAULT_BASE_MODEL,
    adapter: str = DEFAULT_ADAPTER,
    k: int = 4,
    max_new_tokens: int = 320,
    no_generate: bool = False,
) -> str:
    contexts = retrieve(question, index_dir, k)
    if no_generate:
        return json.dumps(
            [{"score": item["score"], "source": item["source"], "text": item["text"][:700]} for item in contexts],
            indent=2,
        )

    prompt = build_prompt(question, contexts)
    generator = load_generator(base_model, adapter)
    output = generator(prompt, max_new_tokens=max_new_tokens, do_sample=False, return_full_text=False)
    answer = output[0]["generated_text"].strip()
    sources = "\n".join(f"- {item['source']} ({item['score']:.3f})" for item in contexts)
    return f"{answer}\n\nSources:\n{sources}"


def main() -> None:
    parser = argparse.ArgumentParser(description="Ask a Qwen RAG assistant about Edge Impulse docs.")
    parser.add_argument("question")
    parser.add_argument("--index-dir", type=Path, default=DEFAULT_INDEX_DIR)
    parser.add_argument("--base-model", default=DEFAULT_BASE_MODEL)
    parser.add_argument("--adapter", default=DEFAULT_ADAPTER)
    parser.add_argument("--k", type=int, default=4)
    parser.add_argument("--max-new-tokens", type=int, default=320)
    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,
            base_model=args.base_model,
            adapter=args.adapter,
            k=args.k,
            max_new_tokens=args.max_new_tokens,
            no_generate=args.no_generate,
        )
    )


if __name__ == "__main__":
    main()