#!/usr/bin/env python3
"""One-shot Kimodo text encoder: texts on stdin, vectors in --out, then exit so the OS reclaims its RAM.

Reads {"texts": [...]} as JSON on stdin. Loads LLM2Vec assembly B on CPU in bfloat16, encodes each text alone as
LLM2VecEncoder does, and writes a MessagePack answer. The torch thread budget is VILI_THREADS.
"""
import json
import os
import resource
import sys
import time
from pathlib import Path

import msgpack
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent))
from text_encoding import ASSEMBLY, EMBEDDING_DIM, LLAMA_BASE, apply_thread_cap, snapshot_path  # noqa: E402

from peft import PeftModel  # noqa: E402
from transformers import AutoConfig, AutoTokenizer  # noqa: E402

from kimodo.model.llm2vec.llm2vec import LLM2Vec  # noqa: E402
from kimodo.model.load_model import TEXT_ENCODER_PRESETS  # noqa: E402


def load_encoder(hf_home, mntp_id, supervised_id, dtype):
    """Assembly B: the Llama base with the MNTP adapter merged in, then the supervised adapter on top.

    Kimodo's own LLM2VecEncoder loads the MNTP repo as the base (assembly A); this assembly is ours
    (motion-pipeline:ruling:encoder-assembly-b). Everything loads from the local snapshot directories, as the
    research's reference vector did: the model's _name_or_path is then a path, not the repo id, so
    LLM2Vec.prepare_for_tokenization does not wrap the text in Llama's chat header, which is what separates B's
    vector from A's. The identity carries the Kimodo commit and this file's sha256, so a changed loader misses.
    """
    llama, mntp, supervised = (snapshot_path(hf_home, repo) for repo in (LLAMA_BASE, mntp_id, supervised_id))
    tokenizer = AutoTokenizer.from_pretrained(llama)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"
    model_class = LLM2Vec._get_model_class(AutoConfig.from_pretrained(llama).__class__.__name__, enable_bidirectional=True)
    model = model_class.from_pretrained(llama, torch_dtype=dtype)
    model = PeftModel.from_pretrained(model, mntp).merge_and_unload()
    model = PeftModel.from_pretrained(model, supervised)
    return LLM2Vec(model=model, tokenizer=tokenizer)


def main() -> int:
    apply_thread_cap(os.environ)
    texts = json.load(sys.stdin)["texts"]
    out = Path(sys.argv[sys.argv.index("--out") + 1])
    preset = TEXT_ENCODER_PRESETS["llm2vec"]["kwargs"]

    started = time.perf_counter()
    model = load_encoder(os.environ["HF_HOME"], preset["base_model_name_or_path"], preset["peft_model_name_or_path"], getattr(torch, preset["dtype"]))
    model.eval()
    for p in model.parameters():
        p.requires_grad = False
    load_seconds = time.perf_counter() - started

    vectors, encode_seconds = [], []
    for text in texts:
        began = time.perf_counter()
        features = model.tokenize([model.prepare_for_tokenization(model._convert_to_str("", text))])
        with torch.no_grad():
            vector = model.forward(features).detach().cpu().to(torch.float32)[0]
        assert vector.shape == (EMBEDDING_DIM,)
        vectors.append(vector.numpy().tobytes())
        encode_seconds.append(time.perf_counter() - began)

    out.write_bytes(msgpack.packb({
        "vectors": vectors,
        "encodeSeconds": encode_seconds,
        "loadSeconds": load_seconds,
        "placement": {"cpuBytes": sum(p.numel() * p.element_size() for p in model.parameters())},
        "peakRssBytes": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024,
        "device": "cpu",
        "assembly": ASSEMBLY,
    }, use_bin_type=True))
    return 0


sys.exit(main())