from typing import Any, Dict, List, override

import numpy as np

from fine_chunker.model.base import BaseChunkingModel


class ONNXChunkingModel(BaseChunkingModel):
    def __init__(self, model_path: str, tokenizer_path: str, device: str = "cpu"):
        try:
            import onnxruntime as ort
        except ImportError:
            raise ImportError(
                "Library 'onnxruntime' is not installed. "
                "Please install it with 'pip install onnxruntime' to use ONNXChunkingModel."
            )
        from tokenizers import Tokenizer

        self.tokenizer = Tokenizer.from_file(tokenizer_path)

        self.session = ort.InferenceSession(
            model_path, providers=self._get_provider(device)
        )

    @override
    def predict(self, input_ids: List[int], threshold: float) -> Dict[str, Any]:
        input_array = np.array([input_ids], dtype=np.int64)

        actual_input_names = [i.name for i in self.session.get_inputs()]

        input_array = np.array([input_ids], dtype=np.int64)
        inputs = {actual_input_names[0]: input_array}

        if "attention_mask" in actual_input_names:
            inputs["attention_mask"] = np.ones(input_array.shape, dtype=np.int64)

        outputs = self.session.run(None, inputs)
        logits = outputs[0]

        max_l = np.max(logits, axis=-1, keepdims=True)
        exp_l = np.exp(logits - max_l)
        probs = exp_l / np.sum(exp_l, axis=-1, keepdims=True)
        predictions = (probs[0, :, 1] >= threshold).astype(int)

        return {
            "logits": logits[0],
            "predictions": predictions,
        }

    @override
    def decode(self, token_ids: List[int]) -> str:
        return self.tokenizer.decode(token_ids, skip_special_tokens=True).strip()

    @override
    def tokenize_and_predict(self, text: str, threshold: float = 0.5) -> Dict[str, Any]:
        encoding = self.tokenizer.encode(text, add_special_tokens=True)
        input_ids = encoding.ids
        result = self.predict(input_ids, threshold)

        return {
            "input_ids": np.array(input_ids, dtype=np.int64),
            "predictions": result["predictions"],
            "logits": result["logits"],
        }

    def _get_provider(self, device: str) -> List[str]:
        providers = []
        if device == "cuda":
            providers.append("CUDAExecutionProvider")
        elif device == "npu":
            providers.append("OpenVINOExecutionProvider")

        providers.append("CPUExecutionProvider")
        return providers

    @override
    def tokenize(self, text: str, **kwargs) -> List[int]:
        add_special = kwargs.get("add_special_tokens", False)
        encoding = self.tokenizer.encode(text, add_special_tokens=add_special)
        return encoding.ids

    @override
    def max_length(self) -> int:
        return getattr(self.tokenizer, "model_max_length", 8192)
