"""Local release loader and optional text -> BGE-M3 -> Scripture Genome CLI.

Read the accompanying custom model code before allowing trust_remote_code.
INT8 is weight-only storage, dequantized to FP32 here, not INT8 execution.
"""

import argparse
import hashlib
import json
from pathlib import Path

import torch
from safetensors.torch import load_file
from transformers import AutoConfig, AutoModel, AutoTokenizer


def load_head(model_dir, precision="fp32", device="cpu"):
    model_dir = Path(model_dir)
    if precision not in {"fp32", "fp16", "int8"}:
        raise ValueError("precision must be fp32, fp16, or int8")
    config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True, local_files_only=True)
    gene_bytes = (model_dir / "gene_ids.txt").read_bytes()
    if hashlib.sha256(gene_bytes).hexdigest() != config.gene_ids_sha256:
        raise ValueError("gene_ids.txt does not match the model configuration")
    genes = gene_bytes.decode("utf-8").splitlines()
    if len(genes) != config.output_dim or len(set(genes)) != len(genes):
        raise ValueError("Invalid gene count/order file")
    if precision == "int8":
        arrays = load_file(str(model_dir / "head_int8.safetensors"))
        model = AutoModel.from_config(config, trust_remote_code=True)
        expected = {f"{prefix}{i}" for i in (0, 1) for prefix in ("w", "b", "scale")}
        if set(arrays) != expected:
            raise ValueError("Unexpected INT8 tensor keys")
        state = {}
        for i, layer in ((0, 0), (1, 2)):
            w, scale, bias = arrays[f"w{i}"], arrays[f"scale{i}"], arrays[f"b{i}"]
            if w.dtype != torch.int8 or scale.shape != (w.shape[0],):
                raise ValueError("Invalid INT8 weights or scales")
            if not torch.isfinite(scale).all() or not (scale > 0).all() or not torch.isfinite(bias).all():
                raise ValueError("Nonfinite weights or invalid quantization scales")
            state[f"head.{layer}.weight"] = w.float() * scale.float()[:, None]
            state[f"head.{layer}.bias"] = bias.float()
        model.load_state_dict(state, strict=True)
    else:
        # FP16 storage is loaded into FP32 arithmetic here for CPU portability.
        # To run FP16 on a compatible device, explicitly cast the head and input.
        options = {"variant": "fp16"} if precision == "fp16" else {}
        model = AutoModel.from_pretrained(
            model_dir, trust_remote_code=True, local_files_only=True,
            use_safetensors=True, dtype=torch.float32, **options,
        )
    return model.to(device).eval(), genes


def encode_texts(texts, head, tokenizer, encoder):
    """No retrieval: returns dense embeddings and 568-feature predictions."""
    if not texts or any(not isinstance(t, str) or not t.strip() for t in texts):
        raise ValueError("Provide a nonempty list of nonempty strings")
    encoded = tokenizer(texts, padding=True, truncation=False, return_tensors="pt")
    if encoded["input_ids"].shape[1] > head.config.max_length:
        raise ValueError(f"Input exceeds {head.config.max_length} tokens; shorten it explicitly")
    encoder_device = next(encoder.parameters()).device
    head_parameter = next(head.parameters())
    encoded = {name: tensor.to(encoder_device) for name, tensor in encoded.items()}
    head.eval()
    encoder.eval()
    with torch.inference_mode():
        hidden = encoder(**encoded).last_hidden_state
        dense = torch.nn.functional.normalize(hidden[:, 0, :].float(), p=2, dim=-1)
        genome = head(embeddings=dense.to(device=head_parameter.device, dtype=head_parameter.dtype)).genome
    return {"dense": dense.cpu(), "genome": genome.float().cpu()}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-dir", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--precision", choices=("fp32", "fp16", "int8"), default="fp32")
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--base-dir", type=Path, help="Optional locally extracted pinned BGE-M3 folder")
    parser.add_argument("--text", action="append", required=True)
    args = parser.parse_args()
    head, genes = load_head(args.model_dir, args.precision, args.device)
    if args.base_dir:
        base, options = str(args.base_dir), {"local_files_only": True}
    else:
        base = head.config.encoder_name
        options = {"revision": head.config.encoder_revision}
    tokenizer = AutoTokenizer.from_pretrained(base, **options)
    encoder = AutoModel.from_pretrained(base, **options).to(args.device).eval()
    result = encode_texts(args.text, head, tokenizer, encoder)
    print(json.dumps({
        "texts": args.text, "dense_shape": list(result["dense"].shape),
        "genome_shape": list(result["genome"].shape),
        "checkpoint_sha256": head.config.checkpoint_sha256,
        "precision": args.precision, "gene_ids": genes,
        "genomes": result["genome"].tolist(),
    }, ensure_ascii=False, allow_nan=False))


if __name__ == "__main__":
    main()