import hashlib
import json
import sys
from pathlib import Path

import numpy as np
from sentence_transformers import SentenceTransformer

SOURCE = Path("chatbot_docs.jsonl")
INDEX = Path("chatbot_embeddings.npz")
MODEL_NAME = "sentence-transformers/multi-qa-MiniLM-L6-cos-v1"


def load_documents(path):
    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]


def text_hash(text):
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def load_existing_index(path):
    if not path.exists():
        return {}
    with np.load(path) as data:
        return {
            str(doc_id): {"hash": str(doc_hash), "embedding": embedding.copy()}
            for doc_id, doc_hash, embedding in zip(
                data["ids"], data["hashes"], data["embeddings"]
            )
        }


def refresh_index(query):
    documents = load_documents(SOURCE)
    previous = load_existing_index(INDEX)
    hashes = {document["id"]: text_hash(document["text"]) for document in documents}
    changed = [
        document
        for document in documents
        if document["id"] not in previous
        or previous[document["id"]]["hash"] != hashes[document["id"]]
    ]

    model = SentenceTransformer(MODEL_NAME)
    refreshed_vectors = {}
    if changed:
        vectors = model.encode_document(
            [document["text"] for document in changed],
            normalize_embeddings=True,
            show_progress_bar=False,
        )
        refreshed_vectors = {
            document["id"]: vector for document, vector in zip(changed, vectors)
        }

    embeddings = [
        refreshed_vectors[document["id"]]
        if document["id"] in refreshed_vectors
        else previous[document["id"]]["embedding"]
        for document in documents
    ]
    matrix = np.vstack(embeddings)
    np.savez(
        INDEX,
        ids=np.array([document["id"] for document in documents]),
        hashes=np.array([hashes[document["id"]] for document in documents]),
        texts=np.array([document["text"] for document in documents]),
        embeddings=matrix,
    )

    query_embedding = model.encode_query(
        query,
        normalize_embeddings=True,
        show_progress_bar=False,
    )
    scores = matrix @ query_embedding
    best = int(np.argmax(scores))

    print(f"documents indexed: {len(documents)}")
    print(
        "documents refreshed: "
        + (", ".join(document["id"] for document in changed) if changed else "none")
    )
    print(f"top result: {documents[best]['id']} score={float(scores[best]):.4f}")
    print(documents[best]["text"])


if __name__ == "__main__":
    refresh_index(sys.argv[1])