"""ONNX inference wrapper for the Argus Nano classifier.

Loads a quantized (or full-precision) ONNX model and its tokenizer,
then exposes ``predict`` / ``predict_batch`` for classification.
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import Optional

import numpy as np


class OnnxModel:
    """Thin wrapper around an ONNX Runtime session + tokenizer."""

    def __init__(
        self,
        model_path: Optional[Path] = None,
        tokenizer_path: Optional[Path] = None,
        max_length: int = 128,
        providers: Optional[list[str]] = None,
        threads: Optional[int] = None,
    ) -> None:
        import onnxruntime as ort

        model_path = model_path or self._default_model_path()
        tokenizer_path = tokenizer_path or self._default_tokenizer_path()

        if not Path(model_path).exists():
            raise FileNotFoundError(f"ONNX model not found: {model_path}")
        if not Path(tokenizer_path).exists():
            raise FileNotFoundError(f"Tokenizer not found: {tokenizer_path}")

        if providers is None:
            available = ort.get_available_providers()
            if "CUDAExecutionProvider" in available:
                providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
            else:
                providers = ["CPUExecutionProvider"]

        # Auto-select model variant: quantized ops run poorly on CUDA
        # (excessive memcpy), so prefer float32 model when GPU is requested.
        if model_path == self._default_quantized_path():
            wants_gpu = "CUDAExecutionProvider" in providers
            float32_path = self._default_float32_path()
            if wants_gpu and float32_path.exists():
                model_path = float32_path

        # Thread control — default to all cores (ONNX RT default).
        # ARGUS_NANO_THREADS env var or explicit threads param can limit it.
        # 0 or None = all cores.
        if threads is None:
            env_threads = os.environ.get("ARGUS_NANO_THREADS")
            threads = int(env_threads) if env_threads else 0
        opts = ort.SessionOptions()
        if threads > 0:
            opts.intra_op_num_threads = threads
            opts.inter_op_num_threads = 1

        self._session = ort.InferenceSession(str(model_path), sess_options=opts, providers=providers)
        self._provider = self._session.get_providers()[0]
        self._tokenizer = self._load_tokenizer(Path(tokenizer_path))
        self._max_length = max_length

    @property
    def provider(self) -> str:
        """The ONNX Runtime execution provider in use."""
        return self._provider

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def predict(self, text: str) -> tuple[int, float]:
        """Classify a single text and return ``(label, confidence)``."""
        input_ids, attention_mask = self._tokenize(text)
        logits = self._session.run(
            None,
            {"input_ids": input_ids, "attention_mask": attention_mask},
        )[0]  # shape (1, num_labels)
        probs = _softmax(logits[0])
        label = int(np.argmax(probs))
        return label, float(probs[label])

    def predict_batch(
        self,
        texts: list[str],
        batch_size: int | None = None,
        progress_fn=None,
    ) -> list[tuple[int, float]]:
        """Classify texts in chunked batches to bound memory usage.

        *batch_size* defaults to 128 on GPU, 64 on CPU.
        *progress_fn*, if provided, is called with (completed, total) after
        each batch so callers can display a progress indicator.
        """
        if not texts:
            return []
        if batch_size is None:
            batch_size = 128 if "CUDA" in self._provider else 64
        results: list[tuple[int, float]] = []
        total = len(texts)
        for start in range(0, total, batch_size):
            chunk = texts[start : start + batch_size]
            input_ids, attention_mask = self._tokenize_batch(chunk)
            logits = self._session.run(
                None,
                {"input_ids": input_ids, "attention_mask": attention_mask},
            )[0]
            for row in logits:
                probs = _softmax(row)
                label = int(np.argmax(probs))
                results.append((label, float(probs[label])))
            if progress_fn is not None:
                progress_fn(min(start + batch_size, total), total)
        return results

    # ------------------------------------------------------------------
    # Tokenizer helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _load_tokenizer(tokenizer_path: Path):
        """Load the tokenizer — prefer fast ``tokenizers`` library."""
        tokenizer_json = tokenizer_path / "tokenizer.json"
        if tokenizer_json.exists():
            from tokenizers import Tokenizer

            return Tokenizer.from_file(str(tokenizer_json))
        # Fallback to HuggingFace transformers (heavier)
        from transformers import AutoTokenizer

        return AutoTokenizer.from_pretrained(str(tokenizer_path), local_files_only=True)

    def _tokenize(self, text: str) -> tuple[np.ndarray, np.ndarray]:
        """Return ``(input_ids, attention_mask)`` as int64 numpy arrays,
        each with shape ``(1, max_length)``."""
        from tokenizers import Tokenizer as FastTokenizer

        if isinstance(self._tokenizer, FastTokenizer):
            self._tokenizer.enable_truncation(max_length=self._max_length)
            self._tokenizer.enable_padding(length=self._max_length, pad_id=1, pad_token="<pad>")
            enc = self._tokenizer.encode(text)
            ids = np.array([enc.ids], dtype=np.int64)
            mask = np.array([enc.attention_mask], dtype=np.int64)
        else:
            # transformers AutoTokenizer fallback
            out = self._tokenizer(
                text,
                truncation=True,
                padding="max_length",
                max_length=self._max_length,
                return_tensors="np",
            )
            ids = out["input_ids"].astype(np.int64)
            mask = out["attention_mask"].astype(np.int64)
        return ids, mask

    def _tokenize_batch(self, texts: list[str]) -> tuple[np.ndarray, np.ndarray]:
        """Batch-tokenize a list of texts. Returns (input_ids, attention_mask)
        each with shape ``(len(texts), max_length)``."""
        from tokenizers import Tokenizer as FastTokenizer

        if isinstance(self._tokenizer, FastTokenizer):
            self._tokenizer.enable_truncation(max_length=self._max_length)
            self._tokenizer.enable_padding(length=self._max_length, pad_id=1, pad_token="<pad>")
            encodings = self._tokenizer.encode_batch(texts)
            ids = np.array([e.ids for e in encodings], dtype=np.int64)
            mask = np.array([e.attention_mask for e in encodings], dtype=np.int64)
        else:
            out = self._tokenizer(
                texts,
                truncation=True,
                padding="max_length",
                max_length=self._max_length,
                return_tensors="np",
            )
            ids = out["input_ids"].astype(np.int64)
            mask = out["attention_mask"].astype(np.int64)
        return ids, mask

    # ------------------------------------------------------------------
    # Default path resolution
    # ------------------------------------------------------------------

    @staticmethod
    def _default_model_path() -> Path:
        env = os.environ.get("ARGUS_NANO_MODEL_PATH")
        if env:
            return Path(env)
        return OnnxModel._default_quantized_path()

    @staticmethod
    def _default_quantized_path() -> Path:
        repo = Path(__file__).resolve().parent.parent.parent
        return repo / "export" / "argus-nano-quantized.onnx"

    @staticmethod
    def _default_float32_path() -> Path:
        repo = Path(__file__).resolve().parent.parent.parent
        return repo / "export" / "argus-nano.onnx"

    @staticmethod
    def _default_tokenizer_path() -> Path:
        env = os.environ.get("ARGUS_NANO_TOKENIZER_PATH")
        if env:
            return Path(env)
        repo = Path(__file__).resolve().parent.parent.parent
        return repo / "export" / "tokenizer"


def _softmax(x: np.ndarray) -> np.ndarray:
    e = np.exp(x - np.max(x))
    return e / e.sum()
