"""CPU inference using only NumPy, Tokenizers and ONNX Runtime.

The host supplies normalized, optional observations. It also owns admission,
async execution and cancellation; this module creates no background workers.
"""
import json
from pathlib import Path

import numpy as np

from .router import Router


class JsonTokenizer:
    """Read the same exported fast tokenizer without loading Transformers."""
    def __init__(self, path):
        from tokenizers import Tokenizer
        path = Path(path)
        self.tokenizer = Tokenizer.from_file(str(path / 'tokenizer.json'))
        config = json.loads((path / 'tokenizer_config.json').read_text())
        pad = config['pad_token']
        if isinstance(pad, dict):
            pad = pad['content']
        pad_id = self.tokenizer.token_to_id(pad)
        if pad_id is None or config.get('padding_side', 'right') != 'right':
            raise ValueError('Expected an exported tokenizer with right padding')
        self.tokenizer.enable_truncation(max_length=384, direction='right')
        self.tokenizer.enable_padding(direction='right', pad_id=pad_id, pad_token=pad)

    def __call__(self, texts, *, padding, truncation, max_length, return_tensors):
        if (padding is not True or truncation is not True
                or max_length != 384 or return_tensors != 'np'):
            raise ValueError('Unsupported tokenization contract')
        values = self.tokenizer.encode_batch(texts, add_special_tokens=True)
        return {'input_ids': np.asarray([r.ids for r in values], dtype=np.int64),
                'attention_mask': np.asarray([r.attention_mask for r in values], dtype=np.int64)}


class OnnxRouter(Router):
    """Drop-in synchronous typed router, with no PyTorch or Studio dependency."""
    def __init__(self, path, *, admission=None):
        super().__init__(path, backend='onnx', device='cpu', admission=admission)

    def _load(self):
        import onnxruntime as ort
        if self.meta.get('format') != 'onnx':
            raise ValueError('Expected an ONNX artifact')
        self.tokenizer = JsonTokenizer(self.path / 'tokenizer')
        options = ort.SessionOptions()
        options.intra_op_num_threads, options.inter_op_num_threads = 2, 1
        options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        options.add_session_config_entry('session.intra_op.allow_spinning', '0')
        options.add_session_config_entry('session.inter_op.allow_spinning', '0')
        self.model = ort.InferenceSession(str(self.path / 'model.onnx'),
            sess_options=options, providers=['CPUExecutionProvider'])