#!/usr/bin/env python3

print("Loading torch/transformers, please wait...", flush=True)

import readline

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

# Model -> (pooling, max tokens). The bge family uses CLS pooling.
MODELS = {
    "BAAI/bge-small-zh-v1.5": ("cls", 512),
    "BAAI/bge-base-zh-v1.5": ("cls", 512),
    "BAAI/bge-large-zh-v1.5": ("cls", 512),
    "BAAI/bge-m3": ("cls", 8192),
}
DEFAULT_MODEL = "BAAI/bge-small-zh-v1.5"
COMMANDS = ["/select", "/sim", "/lsmod", "/help", "/quit", "/exit"]

device = "cuda" if torch.cuda.is_available() else "cpu"
_loaded = {}      # loaded models: name -> (tokenizer, model)
current = None    # currently selected model


def is_cached(name):
    """Whether the model weights exist in the local Hugging Face cache."""
    from huggingface_hub import try_to_load_from_cache

    return any(
        isinstance(try_to_load_from_cache(name, f), str)
        for f in ("model.safetensors", "pytorch_model.bin")
    )


def select_model(name):
    global current
    if name not in MODELS:
        print(f"Unknown model: {name} (use /lsmod to list models)")
        return
    if name not in _loaded:
        print(f"Loading {name} ...", flush=True)
        tok = AutoTokenizer.from_pretrained(name)
        mod = AutoModel.from_pretrained(name).to(device).eval()
        _loaded[name] = (tok, mod)
    current = name
    print(f"Current model: {name}")


def sim(arg):
    """arg is "text1 | text2"; prints their cosine similarity."""
    a, sep, b = arg.partition("|")
    a, b = a.strip(), b.strip()
    if not sep or not a or not b:
        print("Usage: /sim <text1> | <text2>")
        return
    va, vb = encode([a, b])
    # vectors are L2-normalized, so the dot product equals cosine similarity
    print(f"[{current}] similarity: {(va @ vb).item():.4f}")


def lsmod():
    for name in MODELS:
        mark = "*" if name == current else " "
        state = "cached" if is_cached(name) else "not downloaded"
        print(f"{mark} {name}  [{state}]")


@torch.no_grad()
def encode(texts, batch_size=32):
    """Encode with the current model; returns L2-normalized vectors, shape (N, dim)."""
    tokenizer, model = _loaded[current]
    pooling, max_length = MODELS[current]
    outs = []
    for i in range(0, len(texts), batch_size):
        inputs = tokenizer(
            texts[i : i + batch_size],
            padding=True,
            truncation=True,
            max_length=max_length,
            return_tensors="pt",
        ).to(device)
        hidden = model(**inputs).last_hidden_state
        if pooling == "cls":
            emb = hidden[:, 0]
        else:  # mean
            mask = inputs["attention_mask"].unsqueeze(-1).to(hidden.dtype)
            emb = (hidden * mask).sum(1) / mask.sum(1)
        outs.append(F.normalize(emb, p=2, dim=1).cpu())
    return torch.cat(outs)


def completer(text, state):
    """Tab completion: commands at line start, model names after /select."""
    line = readline.get_line_buffer().lstrip()
    if line.startswith("/select "):
        options = [m for m in MODELS if m.startswith(text)]
    elif " " not in line:
        options = [c + " " if c in ("/select", "/sim") else c for c in COMMANDS if c.startswith(text)]
    else:
        options = []
    return options[state] if state < len(options) else None


def setup_readline():
    readline.set_completer_delims(" \t\n")  # do not treat / and - as delimiters
    readline.set_completer(completer)
    doc = readline.__doc__ or ""
    if "libedit" in doc:  # macOS
        readline.parse_and_bind("bind ^I rl_complete")
    else:
        readline.parse_and_bind("tab: complete")


HELP = """Commands:
  /lsmod              list models (* marks the current one)
  /select <model>     switch model (Tab completes)
  /sim <a> | <b>      cosine similarity of two texts (separated by |)
  /help               show this help
  /quit, /exit        quit (empty line or Ctrl+D also works)
  Any other text (including "lsmod" without /) prints its vector.
  To embed text that starts with /, escape it with //, e.g. //lsmod"""


def main():
    import argparse

    parser = argparse.ArgumentParser(description="Enter text, get its embedding vector")
    parser.add_argument("--full", action="store_true", help="print the full vector")
    args = parser.parse_args()

    setup_readline()
    select_model(DEFAULT_MODEL)
    print(HELP)
    while True:
        try:
            text = input("> ").strip()
        except EOFError:
            break
        if not text or text in ("/quit", "/exit"):
            break
        if text.startswith("//"):
            text = text[1:]  # escape: treat text starting with / as plain text
            cmd = None
        else:
            cmd = text if text.startswith("/") else None
        if cmd == "/lsmod":
            lsmod()
        elif cmd == "/help":
            print(HELP)
        elif cmd == "/sim" or (cmd and cmd.startswith("/sim ")):
            sim(cmd[len("/sim"):])
        elif cmd == "/select" or (cmd and cmd.startswith("/select ")):
            name = cmd[len("/select"):].strip()
            if name:
                select_model(name)
            else:
                print("Usage: /select <model>")
        elif cmd:
            print(f"Unknown command: {cmd} (type /help)")
        else:
            vec = encode([text])[0]
            print(f"[{current}] dim: {vec.numel()}  norm: {vec.norm().item():.4f}")
            if args.full:
                print(vec.tolist())
            else:
                print("first 8 values:", [round(x, 4) for x in vec[:8].tolist()])


if __name__ == "__main__":
    main()