"""Commerce query encoder (Korean + English) for EmbeddingGemma 2 product-image indexes.

Search text -> L2-normalized 768-d vector that you compare (dot product) with product-image
vectors produced by google/embeddinggemma-2. The image index does not change.

Two INT8 text encoders share one tokenizer:
  ko/  K4 - trained for Korean product search (77 tokens; long inputs keep the first 57 and last 20 tokens)
  en/  E0 - English product search (64 tokens, plain truncation)
route="hangul" (default) sends a query that contains Hangul to ko/ and every other query to en/.
route="ko" or route="en" uses one encoder for everything.

Weights are row-wise INT8 (`<name>.__scale` holds one fp16 scale per row), dequantized to float32 at load.
"""
import json
import re
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file
from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel

HANGUL = re.compile(r"[가-힣ㄱ-ㆎ]")


def _resolve(path_or_repo, revision=None):
    p = Path(path_or_repo)
    if p.is_dir():
        return p
    from huggingface_hub import snapshot_download
    return Path(snapshot_download(str(path_or_repo), revision=revision))


def head_tail(ids, max_len, tail):
    """Keep the first max_len - tail and the last tail token ids (BOS and EOS survive)."""
    ids = list(ids)
    return ids if len(ids) <= max_len else ids[:max_len - tail] + ids[-tail:]


class TinyQueryEncoder(nn.Module):
    def __init__(self, text_config, out_dim=768, max_tokens=64, tail_tokens=0):
        super().__init__()
        self.enc = CLIPTextModel(text_config)
        self.align = nn.Linear(text_config.hidden_size, out_dim)
        self.max_tokens = min(max_tokens, 77)
        self.tail_tokens = tail_tokens
        self.tok = None

    @classmethod
    def load(cls, folder, tokenizer):
        folder = Path(folder)
        meta = json.loads((folder / "config.json").read_text())
        base = json.loads((folder / "base_config.json").read_text())
        text_cfg = CLIPTextConfig.from_dict(base["text_config"] if base.get("model_type") == "clip" else base)
        tail = (meta.get("truncation") or {}).get("tail_tokens", 0)
        m = cls(text_cfg, out_dim=meta["out_dim"], max_tokens=meta["max_tokens"], tail_tokens=tail)
        raw = load_file(str(folder / "model.safetensors"))
        sd = {k: (v.float() * raw[f"{k}.__scale"].float()[:, None]) if f"{k}.__scale" in raw else v.float()
              for k, v in raw.items() if not k.endswith(".__scale")}
        missing, unexpected = m.load_state_dict(sd, strict=False)
        ref = m.state_dict()
        if unexpected or any(ref[k].dtype.is_floating_point for k in missing):
            raise RuntimeError(f"weight mismatch in {folder}: missing={missing} unexpected={unexpected}")
        m.tok = tokenizer
        return m.eval()

    def _batch(self, texts):
        if not self.tail_tokens:
            return self.tok(texts, padding=True, truncation=True, max_length=self.max_tokens, return_tensors="pt")
        ids = [head_tail(x, self.max_tokens, self.tail_tokens) for x in self.tok(texts, truncation=False)["input_ids"]]
        n = max(len(x) for x in ids)
        a = torch.full((len(ids), n), int(self.tok.pad_token_id or 0), dtype=torch.long)
        m = torch.zeros(len(ids), n, dtype=torch.long)
        for i, x in enumerate(ids):
            a[i, :len(x)] = torch.tensor(x)
            m[i, :len(x)] = 1
        return {"input_ids": a, "attention_mask": m}

    @torch.no_grad()
    def encode(self, texts, batch_size=256):
        dev = next(self.parameters()).device
        out = []
        for i in range(0, len(texts), batch_size):
            b = {k: v.to(dev) for k, v in self._batch(list(texts[i:i + batch_size])).items()}
            out.append(F.normalize(self.align(self.enc(**b).pooler_output).float(), dim=1).cpu().numpy())
        return np.concatenate(out) if out else np.zeros((0, 768), np.float32)


class CommerceQueryEncoder:
    def __init__(self, ko, en, route="hangul"):
        self.ko, self.en, self.route = ko, en, route

    @classmethod
    def from_pretrained(cls, path_or_repo="ThakiCloud/eg2-tiny-query-commerce-ko", revision=None, device="cpu", route="hangul"):
        if route not in ("hangul", "ko", "en"):
            raise ValueError("route must be 'hangul', 'ko' or 'en'")
        root = _resolve(path_or_repo, revision)
        tok = AutoTokenizer.from_pretrained(str(root))
        return cls(TinyQueryEncoder.load(root / "ko", tok).to(device), TinyQueryEncoder.load(root / "en", tok).to(device), route)

    def encode(self, texts, batch_size=256):
        """texts: list[str] -> np.ndarray (n, 768), float32, L2-normalized. No query prefix."""
        if isinstance(texts, str):
            texts = [texts]
        texts = list(texts)
        if self.route != "hangul":
            return (self.ko if self.route == "ko" else self.en).encode(texts, batch_size)
        is_ko = [bool(HANGUL.search(t)) for t in texts]
        out = np.zeros((len(texts), 768), np.float32)
        for flag, model in ((True, self.ko), (False, self.en)):
            idx = [i for i, f in enumerate(is_ko) if f == flag]
            if idx:
                out[idx] = model.encode([texts[i] for i in idx], batch_size)
        return out