#!/usr/bin/env python3

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

import os
import readline

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

# Model -> default parameters. The bge family uses CLS pooling.
#   pooling:    "cls" (first token) or "mean" (average of all tokens)
#   max_length: max tokens per text; longer input is truncated
#   batch_size: texts encoded per forward pass (speed/memory only, no effect on vectors)
MODELS = {
    "BAAI/bge-small-zh-v1.5": {"pooling": "cls", "max_length": 512, "batch_size": 32},
    "BAAI/bge-base-zh-v1.5": {"pooling": "cls", "max_length": 512, "batch_size": 32},
    "BAAI/bge-large-zh-v1.5": {"pooling": "cls", "max_length": 512, "batch_size": 32},
    "BAAI/bge-m3": {"pooling": "cls", "max_length": 8192, "batch_size": 32},
}
PARAM_KEYS = ["pooling", "max_length", "batch_size"]
DEFAULT_MODEL = "BAAI/bge-small-zh-v1.5"
LAST_MODEL_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".cache", "last_model")
COMMANDS = ["/select", "/sim", "/tokens", "/params", "/set", "/reset", "/lsmod", "/help", "/quit", "/exit"]

device = "cuda" if torch.cuda.is_available() else "cpu"
_loaded = {}      # loaded models: name -> (tokenizer, model)
_params = {name: dict(d) for name, d in MODELS.items()}  # live, per-model parameters
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 load_last_model():
    """The model selected last time, or DEFAULT_MODEL if unknown."""
    try:
        with open(LAST_MODEL_FILE) as f:
            name = f.read().strip()
    except OSError:
        return DEFAULT_MODEL
    return name if name in MODELS else DEFAULT_MODEL


def save_last_model(name):
    try:
        os.makedirs(os.path.dirname(LAST_MODEL_FILE), exist_ok=True)
        with open(LAST_MODEL_FILE, "w") as f:
            f.write(name)
    except OSError:
        pass


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
    save_last_model(name)
    print(f"Current model: {name}")


def params():
    for k in PARAM_KEYS:
        v, d = _params[current][k], MODELS[current][k]
        print(f"  {k:<11} = {v}" + ("" if v == d else f"   (default: {d})"))


def set_param(arg):
    """arg is "<key> <value>"; validates and applies it to the current model."""
    parts = arg.split()
    if len(parts) != 2 or parts[0] not in PARAM_KEYS:
        print("Usage: /set <pooling|max_length|batch_size> <value>")
        return
    key, val = parts
    if key == "pooling":
        if val not in ("cls", "mean"):
            print("pooling must be 'cls' or 'mean'")
            return
    else:
        try:
            val = int(val)
        except ValueError:
            print(f"{key} must be an integer")
            return
        limit = _loaded[current][0].model_max_length if key == "max_length" else 4096
        if key == "max_length" and limit > 100_000:  # tokenizer reports no real limit
            limit = 8192
        if not 1 <= val <= limit:
            print(f"{key} must be between 1 and {limit}")
            return
    _params[current][key] = val
    print(f"[{current}] {key} = {val}")


def reset_params():
    _params[current] = dict(MODELS[current])
    print(f"[{current}] parameters reset to defaults")


def show_tokens(text, preview=50):
    """Show how the current model's tokenizer splits text."""
    text = text.strip()
    if not text:
        print("Usage: /tokens <text>")
        return
    tokenizer = _loaded[current][0]
    ids = tokenizer(text, truncation=False)["input_ids"]  # includes special tokens
    toks = tokenizer.convert_ids_to_tokens(ids)
    shown = ", ".join(f"{t}({i})" for t, i in zip(toks[:preview], ids[:preview]))
    more = f", ... (+{len(ids) - preview} more)" if len(ids) > preview else ""
    print(f"[{current}] {len(ids)} tokens: {shown}{more}")
    max_length = _params[current]["max_length"]
    if len(ids) > max_length:
        print(f"  note: exceeds max_length={max_length}; "
              f"the last {len(ids) - max_length} tokens will be truncated")


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):
    """Encode with the current model; returns L2-normalized vectors, shape (N, dim)."""
    tokenizer, model = _loaded[current]
    p = _params[current]
    pooling, max_length, batch_size = p["pooling"], p["max_length"], p["batch_size"]
    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 line.startswith("/set "):
        words = line.split(" ")
        if len(words) <= 2:  # completing the key
            options = [k + " " for k in PARAM_KEYS if k.startswith(text)]
        elif len(words) == 3 and words[1] == "pooling":
            options = [v for v in ("cls", "mean") if v.startswith(text)]
        else:
            options = []
    elif " " not in line:
        options = [c + " " if c in ("/select", "/sim", "/tokens", "/set") 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 |)
  /tokens <text>      show how the text is split into tokens
  /params             show the current model's parameters
  /set <key> <value>  change a parameter: pooling (cls|mean), max_length, batch_size
  /reset              restore the current model's default parameters
  /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(load_last_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 == "/tokens" or (cmd and cmd.startswith("/tokens ")):
            show_tokens(cmd[len("/tokens"):])
        elif cmd == "/params":
            params()
        elif cmd == "/reset":
            reset_params()
        elif cmd == "/set" or (cmd and cmd.startswith("/set ")):
            set_param(cmd[len("/set"):])
        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()