"""中文嵌入：CPU查表、NPU编码器、长文重叠窗口加权聚合。"""
import os
from pathlib import Path
import numpy as np
from tokenizers import Tokenizer
from rknnlite.api import RKNNLite
import rknnlite.api.rknn_lite as lite_module
from rknnlite.api.rknn_runtime import RKNNRuntime

INSTRUCTION = '为这个句子生成表示以用于检索相关文章：'

def normalize(value):
    value = np.asarray(value, dtype=np.float32).reshape(-1)
    norm = np.linalg.norm(value)
    if value.shape != (512,) or not np.isfinite(value).all() or norm < 1e-8:
        raise RuntimeError('无效嵌入输出')
    return value / norm

class Engine:
    def __init__(self, directory):
        self.directory = Path(directory)
        self.tokenizer = Tokenizer.from_file(str(self.directory / 'tokenizer.json'))
        self.tokenizer.no_truncation()
        self.tokenizer.no_padding()
        with np.load(self.directory / 'embedding-tables.npz') as data:
            self.tables = {key: data[key] for key in data.files}
        directory = self.directory
        class LocalRuntime(RKNNRuntime):
            def _get_rknn_api_lib_path(self):
                return str(directory / 'librknnrt.so')
        lite_module.RKNNRuntime = LocalRuntime
        self.runtimes = {}
        try:
            for length in [128, 512]:
                runtime = RKNNLite(verbose=False)
                if runtime.load_rknn(str(directory / f'encoder-{length}.rknn')) != 0 or runtime.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2) != 0:
                    runtime.release()
                    raise RuntimeError('初始化NPU失败')
                self.runtimes[length] = runtime
        except Exception:
            self.release()
            raise
    def release(self):
        for runtime in self.runtimes.values():
            runtime.release()
        self.runtimes.clear()
    def encode_window(self, ids):
        count = len(ids) + 2
        length = 128 if count <= 128 else 512
        padded = np.zeros((1, length), dtype=np.int64)
        padded[0, :count] = [101, *ids, 102]
        t = self.tables
        x = t['word'][padded] + t['position'][np.arange(length)][None, :, :] + t['token'][0]
        mean = x.mean(axis=-1, keepdims=True)
        var = ((x - mean) ** 2).mean(axis=-1, keepdims=True)
        x = (x - mean) / np.sqrt(var + t['eps']) * t['weight'] + t['bias']
        mask = np.full((1, 1, length, 1), -10000.0, dtype=np.float32)
        mask[:, :, :count, :] = 0
        output = self.runtimes[length].inference(inputs=[np.ascontiguousarray(x, dtype=np.float32), mask], data_format=['nchw', 'nhwc'])
        if output is None:
            raise RuntimeError('NPU推理失败')
        return normalize(output[0])
    def embed(self, text, kind='auto'):
        if not isinstance(text, str) or not text.strip():
            raise ValueError('文本不能为空')
        # Cloudreve4.19.1固定文档模板，可区分文档与查询，不改其源码。
        is_document = kind == 'document' or (kind == 'auto' and text.startswith('Chunk #') and " in a file named '" in text)
        if kind == 'query' or not is_document:
            text = INSTRUCTION + text
        ids = self.tokenizer.encode(text, add_special_tokens=False).ids
        windows = [ids[start:start + 510] for start in range(0, len(ids), 494)]
        if len(windows) > int(os.environ.get('MAX_WINDOWS', '32')):
            raise ValueError('文本窗口数量超限，请减小分块')
        vectors = [self.encode_window(window) for window in windows]
        # 所有窗口均参与聚合，不截掉尾部；单向量兼容Meilisearch REST。
        weights = np.array([len(window) for window in windows], dtype=np.float32)
        result = normalize(np.average(vectors, axis=0, weights=weights))
        return result.tolist(), len(ids), len(windows)