#!/usr/bin/env python3

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

import os

import torch
import torch.nn.functional as F
from prompt_toolkit import PromptSession
from prompt_toolkit.application import get_app
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.completion import Completer, Completion
from prompt_toolkit.filters import Condition
from prompt_toolkit.history import FileHistory
from prompt_toolkit.key_binding import KeyBindings
from transformers import AutoModel, AutoTokenizer

# Model -> default parameters. The bge family uses CLS pooling, Qwen3-Embedding uses last-token pooling.
#   pooling:    "cls" (first token), "mean" (average of all tokens) or "last" (last token)
#   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},
    "Qwen/Qwen3-Embedding-8B": {"pooling": "last", "max_length": 8192, "batch_size": 4},
}
# Extra loading options. Qwen3-Embedding is a decoder: left padding puts the last real token at
# position -1, and bfloat16 halves the memory of the 8B weights.
TOKENIZER_OPTS = {"Qwen/Qwen3-Embedding-8B": {"padding_side": "left"}}
MODEL_OPTS = {"Qwen/Qwen3-Embedding-8B": {"torch_dtype": torch.bfloat16}}
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")
HISTORY_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".cache", "history")
COMMANDS = ["/select", "/sim", "/tokens", "/params", "/set", "/reset", "/lsmod", "/help", "/quit", "/exit"]

if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():  # Apple Silicon GPU
    device = "mps"
else:
    device = "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, **TOKENIZER_OPTS.get(name, {}))
        mod = AutoModel.from_pretrained(name, **MODEL_OPTS.get(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", "last"):
            print("pooling must be 'cls', 'mean' or 'last'")
            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]
        elif pooling == "last":
            if tokenizer.padding_side == "left":
                emb = hidden[:, -1]
            else:  # right padding: pick the last non-pad token of each row
                last = inputs["attention_mask"].sum(1) - 1
                emb = hidden[torch.arange(hidden.size(0), device=hidden.device), last]
        else:  # mean
            mask = inputs["attention_mask"].unsqueeze(-1).to(hidden.dtype)
            emb = (hidden * mask).sum(1) / mask.sum(1)
        outs.append(F.normalize(emb.float(), p=2, dim=1).cpu())
    return torch.cat(outs)


def candidates(line, text):
    """Completions for the word `text` at the end of `line`: commands at line start, model names after /select."""
    line = line.lstrip()
    if line.startswith("/select "):
        return [m for m in MODELS if m.startswith(text)]
    if line.startswith("/set "):
        words = line.split(" ")
        if len(words) <= 2:  # completing the key
            return [k + " " for k in PARAM_KEYS if k.startswith(text)]
        if len(words) == 3 and words[1] == "pooling":
            return [v for v in ("cls", "mean", "last") if v.startswith(text)]
        return []
    if " " not in line:
        return [c + " " if c in ("/select", "/sim", "/tokens", "/set") else c for c in COMMANDS if c.startswith(text)]
    return []


class InputCompleter(Completer):
    """Completes while typing."""

    def get_completions(self, document, complete_event):
        text = document.get_word_before_cursor(WORD=True)
        for option in candidates(document.text_before_cursor, text):
            yield Completion(option, start_position=-len(text))


def make_session():
    try:
        os.makedirs(os.path.dirname(HISTORY_FILE), exist_ok=True)
    except OSError:
        pass
    bindings = KeyBindings()

    @Condition
    def completion_selected():
        state = get_app().current_buffer.complete_state
        return state is not None and state.current_completion is not None

    @bindings.add("enter", filter=completion_selected)
    def _(event):
        # Enter on a highlighted candidate accepts it (already inserted) instead of submitting
        event.current_buffer.complete_state = None

    return PromptSession(
        key_bindings=bindings,
        completer=InputCompleter(),
        complete_while_typing=True,
        auto_suggest=AutoSuggestFromHistory(),
        history=FileHistory(HISTORY_FILE),
    )


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|last), 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()

    session = make_session()
    select_model(load_last_model())
    print(HELP)
    while True:
        try:
            text = session.prompt("> ").strip()
        except (EOFError, KeyboardInterrupt):
            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()