"""Integrate multilingual depth classification with Studio routing."""
import json
from pathlib import Path

from studio.core.utils.reply_modes.encoder_depth import EncoderDepthRouter

from .features import NETWORK_TYPE, contract
from .history import history_view


class VoiceDepthRouter(EncoderDepthRouter):
    def __init__(self, model, *, device='cpu', backend='onnx', allow_prototype=False):
        metadata = json.loads((Path(model).expanduser()/'depth_encoder.json').read_text())
        self.history_policy = metadata.get('history_policy', 'first80')
        if (metadata.get('network_type') != NETWORK_TYPE
                or metadata.get('observation_contract') != contract(self.history_policy)):
            raise ValueError('Incompatible multilingual voice checkpoint')
        if metadata.get('pooling') not in ('cls', 'mean'):
            raise ValueError('Invalid pooling')
        if metadata.get('threshold', .5) != .5:
            raise ValueError('The routing contract requires threshold=0.5')
        if metadata.get('report', {}).get('trained_audio_state') is True:
            if metadata['report'].get('measured_synthetic_speech_training_rows', 0) <= 0:
                raise ValueError('Voice training requires measured observation provenance')
        super().__init__(model, device=device, backend=backend, allow_prototype=allow_prototype)

    def _input_key(self, text, history, audio_state):
        return super()._input_key(text, history_view(history, self.history_policy), audio_state)

    def classify(self, text, memory_prefetch_hint=False, history=None, *, audio_state=None):
        return super().classify(text, memory_prefetch_hint,
                                history_view(history, self.history_policy), audio_state=audio_state)

    def _load(self):
        if self.backend == 'onnx':
            from transformers import AutoTokenizer
            import onnxruntime as ort
            from voicemem.engine import cpu_session_options
            self._tokenizer = AutoTokenizer.from_pretrained(self.path/'tokenizer', local_files_only=True,
                                                           trust_remote_code=False, fix_mistral_regex=False)
            self._model = ort.InferenceSession(str(self.path/'model.onnx'),
                sess_options=cpu_session_options(ort), providers=['CPUExecutionProvider'])
        else:
            from .model import load
            self._model, self._tokenizer = load(self.path, self.metadata, self.device)


def create(*, model, base_url='', api_key='', options=None):
    options = dict(options or {})
    if base_url or api_key or set(options)-{'device', 'backend', 'allow_prototype'}:
        raise ValueError('Voice router accepts local runtime options only')
    if 'allow_prototype' in options and type(options['allow_prototype']) is not bool:
        raise ValueError('allow_prototype must be a boolean')
    if 'device' in options and (not isinstance(options['device'], str) or not options['device']):
        raise ValueError('device must be a nonempty device name')
    return VoiceDepthRouter(model, **options)