"""Load pinned BioLORD weights and cache normalized text embeddings."""

import argparse
import json
import os
from pathlib import Path

import torch

from .provenance import TEXT_ID, TEXT_REVISION, TEXT_IDENTITY, VERSION, sha256


class TextEncoder:
    """Load frozen BioLORD on demand and cache normalized sentence vectors on CPU."""

    def __init__(self, cache=None, device="cpu", model_path=None):
        self.device, self.model_path, self.model = device, model_path, None
        self.embeddings = {}
        if cache is not None and Path(cache).exists():
            payload = torch.load(cache, map_location="cpu", weights_only=True)
            if payload.get("identity") != TEXT_IDENTITY:
                raise ValueError("Text bank does not match pinned BioLORD processing")
            for text, vector in payload["embeddings"].items():
                if vector.shape != (768,) or not torch.isfinite(vector).all():
                    raise ValueError("Invalid text bank vector")
                self.embeddings[text] = vector.float()

    def encode(self, strings):
        """Return separate [N,768] CPU vectors; reject texts above the token limit."""
        if not isinstance(strings, list) or not strings or any(not isinstance(s, str) or not s.strip() for s in strings):
            raise ValueError("Encode a nonempty list of separate, nonempty text strings")
        missing = list(dict.fromkeys(s for s in strings if s not in self.embeddings))
        if missing:
            if self.model is None:
                from sentence_transformers import SentenceTransformer

                local = self.model_path or (
                    Path("third-party/biolord")
                    if Path("third-party/biolord").exists()
                    else None
                )
                if local is not None:
                    local = Path(local)
                    source = json.loads((local / "source.json").read_text())
                    if (
                        source.get("model_id") != TEXT_ID
                        or source.get("revision") != TEXT_REVISION
                    ):
                        raise ValueError(
                            "Local text weights do not match pinned BioLORD"
                        )
                    for name, expected in source["files"].items():
                        path = (local / name.replace("\\", "/")).resolve()
                        if (
                            not path.is_relative_to(local.resolve())
                            or sha256(path) != expected
                        ):
                            raise ValueError(
                                "Local text weight/config checksum changed"
                            )
                self.model_path = local
                self.model = SentenceTransformer(
                    str(local) if local else TEXT_ID,
                    revision=None if local else TEXT_REVISION,
                    device=self.device,
                    trust_remote_code=False,
                    cache_folder=os.environ.get("HF_HOME", "cache/huggingface"),
                )
                self.model.eval().requires_grad_(False)
            for text in missing:
                if (
                    len(self.model.tokenizer.encode(text, add_special_tokens=True))
                    > self.model.max_seq_length
                ):
                    raise ValueError(
                        "Text exceeds BioLORD context; do not silently truncate decisive wording"
                    )
            with torch.inference_mode():
                vectors = self.model.encode(
                    missing,
                    normalize_embeddings=True,
                    convert_to_tensor=True,
                    show_progress_bar=False,
                    batch_size=32,
                )
            for text, vector in zip(missing, vectors):
                if vector.shape != (768,) or not torch.isfinite(vector).all():
                    raise ValueError("Text encoder produced an invalid vector")
                self.embeddings[text] = vector.detach().cpu().float().clone()
        return torch.stack([self.embeddings[s] for s in strings])

    def save(self, path):
        """Save cached vectors with their encoder identity."""
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        torch.save({"identity": TEXT_IDENTITY, "embeddings": self.embeddings}, path)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--texts", type=Path, required=True)
    parser.add_argument("--output", type=Path, default=Path(f"cache/{VERSION}/text.pt"))
    parser.add_argument("--device", default="cpu")
    args = parser.parse_args()
    encoder = TextEncoder(args.output, args.device)
    encoder.encode(json.loads(args.texts.read_text()))
    encoder.save(args.output)
    print(json.dumps({"texts": len(encoder.embeddings), "identity": TEXT_IDENTITY}))


if __name__ == "__main__":
    main()