"""Load a checkpoint or complete GGUF with a PyTorch or native MLX backend.

Q8_0 is dequantized on load; this is not a native quantized GGML backend.
"""
import tempfile
from pathlib import Path
import laya
from .gguf_io import restore_gguf


class Model:
    def __init__(self, path, device=None, *, backend='torch', dtype=None, batch_size=16):
        if backend not in ('torch', 'mlx'):
            raise ValueError('backend must be torch or mlx')
        self._temporary = None
        if backend == 'mlx':
            from .mlx_runtime import MLXModel
            self._model = MLXModel(path, device, dtype=dtype or 'float16', batch_size=batch_size)
            return
        if dtype is not None:
            raise ValueError('dtype is only supported by the MLX entry point')
        try:
            restored = Path(path)
            if not restored.is_dir():
                self._temporary = tempfile.TemporaryDirectory(prefix='xdecision-')
                restored = Path(self._temporary.name)/'model'
                restore_gguf(path, None, restored)
            self._model = laya.load(str(restored), device=device or 'cpu')
        except BaseException:
            if self._temporary is not None:
                self._temporary.cleanup()
            raise

    def predict(self, state, questions):
        if self._model is None:
            raise RuntimeError('Model is closed')
        result = self._model.predict(state, questions)
        result['model'] = 'xDecision'
        return result

    def close(self):
        if hasattr(self._model, 'close'):
            self._model.close()
        self._model = None
        if self._temporary is not None:
            self._temporary.cleanup()

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.close()


def load(path, device=None, *, backend='torch', dtype=None, batch_size=16):
    return Model(path, device, backend=backend, dtype=dtype, batch_size=batch_size)