"""
scripts/model.py
----------------
Retrieval-Augmented Generation pipeline supporting three retrieval strategies:

  dense  — sentence-transformers embeddings + FAISS  (deep learning)
  tfidf  — TF-IDF bag-of-words + cosine similarity   (classical ML)
  random — random chunk selection                     (naive baseline)

Usage (from project root):
    .venv/bin/python scripts/model.py catan 'What resources do you need to build a settlement?'
    .venv/bin/python scripts/model.py monopoly 'How much money does each player start with?' --retriever tfidf
    .venv/bin/python scripts/model.py catan --retriever random

Requires scripts/build_features.py to have been run first for the given game.
Requires OPENAI_API_KEY to be set in the environment.
"""

import json
import os
import sys

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

DEFAULT_K         = 3
RERANK_CANDIDATES = 10
RERANK_TOP_N      = 3
DEFAULT_RETRIEVER = "dense"
RETRIEVERS        = ["dense", "tfidf", "random"]
OPENAI_MODEL      = "gpt-4o-mini"
CROSS_ENCODER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"

# ---------------------------------------------------------------------------
# Per-game singleton registry
# ---------------------------------------------------------------------------

_registry: dict[str, dict] = {}   # game -> loaded artifacts


def _make_system_prompt(game: str) -> str:
    game_name = game.title()
    return (
        f"You are an expert on the board game {game_name}. "
        "Answer the user's question using ONLY the rulebook sections provided as context.\n\n"
        "Rules:\n"
        "1. Base your answer SOLELY on the provided context. Do not use outside knowledge.\n"
        "2. If the answer can be inferred or directly found in the context, state it clearly "
        "and cite the source section, e.g. \"According to [Almanac: ROBBER]...\".\n"
        "3. Only say \"The rules don't specify this.\" if NO relevant information appears in ANY "
        "of the retrieved sections. Do not abstain if the answer is present or can be reasonably "
        "inferred from the context.\n"
        "4. Do not invent or extrapolate rules beyond what the text says.\n"
        "5. Be concise and precise."
    )


def _load(game: str):
    """Lazy-load and cache all retrieval artifacts for the given game."""
    if game in _registry:
        return

    try:
        import faiss
        from sentence_transformers import SentenceTransformer
        from openai import OpenAI
    except ImportError as e:
        print(f"Missing dependency: {e}\nRun: .venv/bin/pip install -r requirements.txt", file=sys.stderr)
        sys.exit(1)

    chunks_path = f"models/{game}/chunks.json"
    index_path  = f"models/{game}/faiss.index"

    if not os.path.exists(chunks_path) or not os.path.exists(index_path):
        print(
            f"Index files not found for '{game}'. "
            f"Run: .venv/bin/python scripts/build_features.py {game}",
            file=sys.stderr,
        )
        sys.exit(1)

    if not os.environ.get("OPENAI_API_KEY"):
        print("OPENAI_API_KEY is not set.", file=sys.stderr)
        sys.exit(1)

    print(f"Loading embedding model (all-MiniLM-L6-v2)...")
    dense_model = SentenceTransformer("all-MiniLM-L6-v2")

    print(f"Loading FAISS index for '{game}'...")
    index = faiss.read_index(index_path)

    with open(chunks_path) as f:
        chunks = json.load(f)

    entry: dict = {
        "dense_model":   dense_model,
        "index":         index,
        "chunks":        chunks,
        "client":        OpenAI(),
        "system_prompt": _make_system_prompt(game),
    }

    # TF-IDF artifacts (classical ML) — optional, built by build_features.py
    tfidf_matrix_path     = f"models/{game}/tfidf_matrix.npz"
    tfidf_vectorizer_path = f"models/{game}/tfidf_vectorizer.pkl"
    if os.path.exists(tfidf_matrix_path) and os.path.exists(tfidf_vectorizer_path):
        import pickle
        import scipy.sparse
        entry["tfidf_matrix"]     = scipy.sparse.load_npz(tfidf_matrix_path)
        with open(tfidf_vectorizer_path, "rb") as f:
            entry["tfidf_vectorizer"] = pickle.load(f)
        print(f"Loaded TF-IDF index for '{game}'.")

    # Cross-encoder for reranking (loaded once, shared across games)
    if "_cross_encoder" not in _registry:
        from sentence_transformers import CrossEncoder
        print(f"Loading cross-encoder model ({CROSS_ENCODER_MODEL})...")
        _registry["_cross_encoder"] = CrossEncoder(CROSS_ENCODER_MODEL)

    _registry[game] = entry
    print(f"Ready — {len(chunks)} chunks loaded for '{game}'.\n")


# ---------------------------------------------------------------------------
# Retrieval — three strategies
# ---------------------------------------------------------------------------

def _dense_retrieve(query: str, game: str, k: int) -> list[dict]:
    """Deep learning: sentence-transformer embeddings + FAISS cosine search."""
    import faiss
    reg = _registry[game]

    query_emb = reg["dense_model"].encode([query], convert_to_numpy=True)
    faiss.normalize_L2(query_emb)
    scores, indices = reg["index"].search(query_emb, k)

    results = []
    for score, idx in zip(scores[0], indices[0]):
        if idx >= 0:
            chunk = reg["chunks"][idx].copy()
            chunk["retrieval_score"] = float(score)
            chunk["chunk_idx"]       = int(idx)
            results.append(chunk)
    return results


def _tfidf_retrieve(query: str, game: str, k: int) -> list[dict]:
    """Classical ML: TF-IDF bag-of-words + cosine similarity."""
    from sklearn.metrics.pairwise import cosine_similarity
    reg = _registry[game]

    if "tfidf_vectorizer" not in reg:
        print(
            f"TF-IDF index not found for '{game}'. "
            f"Re-run: .venv/bin/python scripts/build_features.py {game}",
            file=sys.stderr,
        )
        return []

    query_vec = reg["tfidf_vectorizer"].transform([query])
    scores    = cosine_similarity(query_vec, reg["tfidf_matrix"])[0]
    top_idx   = scores.argsort()[::-1][:k]

    results = []
    for idx in top_idx:
        chunk = reg["chunks"][idx].copy()
        chunk["retrieval_score"] = float(scores[idx])
        chunk["chunk_idx"]       = int(idx)
        results.append(chunk)
    return results


def _random_retrieve(query: str, game: str, k: int) -> list[dict]:
    """Naive baseline: randomly select k chunks (no query signal)."""
    import random
    reg     = _registry[game]
    chunks  = reg["chunks"]
    indices = random.sample(range(len(chunks)), min(k, len(chunks)))

    results = []
    for idx in indices:
        chunk = chunks[idx].copy()
        chunk["retrieval_score"] = 0.0
        chunk["chunk_idx"]       = int(idx)
        results.append(chunk)
    return results


def _rerank(query: str, candidates: list[dict], top_n: int = RERANK_TOP_N) -> list[dict]:
    """Re-score candidate chunks with a cross-encoder and return the top-n."""
    cross_encoder = _registry["_cross_encoder"]
    pairs = [[query, c["text"]] for c in candidates]
    scores = cross_encoder.predict(pairs)

    for chunk, score in zip(candidates, scores):
        chunk["rerank_score"] = float(score)

    ranked = sorted(candidates, key=lambda c: c["rerank_score"], reverse=True)
    return ranked[:top_n]


def retrieve(
    query:     str,
    game:      str,
    k:         int = DEFAULT_K,
    retriever: str = DEFAULT_RETRIEVER,
    rerank:    bool = False,
) -> list[dict]:
    """Return the top-k most relevant chunks using the specified retrieval strategy.

    If rerank=True, first retrieve RERANK_CANDIDATES chunks, then re-score them
    with a cross-encoder and return the top RERANK_TOP_N.
    """
    _load(game)
    fetch_k = RERANK_CANDIDATES if (rerank and retriever != "random") else k

    if retriever == "tfidf":
        results = _tfidf_retrieve(query, game, fetch_k)
    elif retriever == "random":
        results = _random_retrieve(query, game, k)
    else:
        results = _dense_retrieve(query, game, fetch_k)

    if rerank and retriever != "random":
        results = _rerank(query, results, top_n=RERANK_TOP_N)

    return results


# ---------------------------------------------------------------------------
# Generation (shared across all retrievers)
# ---------------------------------------------------------------------------

def generate(query: str, context_chunks: list[dict], game: str) -> str:
    """Call OpenAI with retrieved context and return a grounded answer."""
    _load(game)
    reg = _registry[game]

    context_str  = "\n\n---\n\n".join(f"[{c['title']}]\n{c['text']}" for c in context_chunks)
    user_message = f"Rulebook context:\n\n{context_str}\n\n---\n\nQuestion: {query}"

    response = reg["client"].chat.completions.create(
        model=OPENAI_MODEL,
        messages=[
            {"role": "system", "content": reg["system_prompt"]},
            {"role": "user",   "content": user_message},
        ],
        temperature=0,
        max_tokens=600,
    )
    return response.choices[0].message.content.strip()


# ---------------------------------------------------------------------------
# Full pipeline
# ---------------------------------------------------------------------------

def query_rag(
    question:  str,
    game:      str,
    k:         int = DEFAULT_K,
    retriever: str = DEFAULT_RETRIEVER,
    rerank:    bool = False,
    verbose:   bool = False,
) -> dict:
    """
    Full RAG pipeline: retrieve top-k chunks, then generate a grounded answer.

    retriever options:
      "dense"  — sentence-transformer embeddings + FAISS  (deep learning)
      "tfidf"  — TF-IDF + cosine similarity               (classical ML)
      "random" — random chunk selection                    (naive baseline)

    If rerank=True, retrieves 10 candidates and re-scores with a cross-encoder
    to select the top 3.
    """
    retrieved = retrieve(question, game, k, retriever, rerank=rerank)
    answer    = generate(question, retrieved, game)

    result = {
        "question":  question,
        "answer":    answer,
        "retriever": retriever,
        "retrieved": [
            {"title": c["title"], "score": c["retrieval_score"], "chunk_idx": c["chunk_idx"]}
            for c in retrieved
        ],
    }

    if verbose:
        print(f"Retriever: {retriever}")
        print(f"Question: {question}\n")
        print("Retrieved chunks:")
        for c in result["retrieved"]:
            print(f"  [{c['score']:.3f}] {c['title']}")
        print(f"\nAnswer:\n{answer}")
        print("\n" + "-" * 60)

    return result


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Query the board game RAG pipeline.")
    parser.add_argument("game",      help="Game to query (e.g. catan, monopoly)")
    parser.add_argument("question",  nargs="?", default="", help="Question to ask")
    parser.add_argument("--retriever", choices=RETRIEVERS, default=DEFAULT_RETRIEVER,
                        help="Retrieval strategy (default: dense)")
    parser.add_argument("--rerank", action="store_true",
                        help="Re-rank candidates with cross-encoder (top 10 → top 3)")
    args = parser.parse_args()

    _defaults = {
        "catan":    "What resources do you need to build a settlement?",
        "monopoly": "How much money does each player start with?",
    }
    question = args.question or _defaults.get(args.game, "")
    query_rag(question, game=args.game, retriever=args.retriever, rerank=args.rerank, verbose=True)