"""OpenVINO LLM provider — Intel hardware inference via optimum-intel."""

from __future__ import annotations

import logging
from collections.abc import AsyncIterator
from typing import Any

from redhat_oracle.core.llm.base import LLMProvider

logger = logging.getLogger(__name__)


class OpenVINOUnavailableError(Exception):
    """Raised when OpenVINO runtime or model is not available."""


class OpenVINOLLM(LLMProvider):
    """OpenVINO LLM provider for Intel NPU/iGPU/CPU inference.

    Uses optimum-intel to load and run quantized models locally.
    """

    def __init__(
        self,
        model_id: str = "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        device: str = "CPU",
    ) -> None:
        self._model_id = model_id
        self._device = device
        self._model: Any = None
        self._tokenizer: Any = None

    def _load_model(self) -> None:
        """Lazily load the OpenVINO model and tokenizer."""
        if self._model is not None:
            return

        try:
            from optimum.intel import OVModelForCausalLM
            from transformers import AutoTokenizer

            logger.info(
                "Loading OpenVINO model: %s on %s",
                self._model_id,
                self._device,
            )
            self._tokenizer = AutoTokenizer.from_pretrained(self._model_id)
            self._model = OVModelForCausalLM.from_pretrained(
                self._model_id,
                export=True,
                device=self._device,
            )
        except ImportError as e:
            raise OpenVINOUnavailableError(
                "optimum-intel not installed. Install with: "
                "pip install optimum[openvino]"
            ) from e
        except Exception as e:
            raise OpenVINOUnavailableError(
                f"Failed to load OpenVINO model: {e}"
            ) from e

    async def _check_available(self) -> None:
        """Check if OpenVINO runtime is available."""
        try:
            import optimum.intel  # noqa: F401
        except ImportError as e:
            raise OpenVINOUnavailableError(
                "optimum-intel not installed"
            ) from e

    async def generate(self, prompt: str, system: str | None = None) -> str:
        """Generate a complete response using OpenVINO."""
        self._load_model()

        full_prompt = prompt
        if system:
            full_prompt = f"{system}\n\n{prompt}"

        inputs = self._tokenizer(full_prompt, return_tensors="pt")
        outputs = self._model.generate(
            **inputs,
            max_new_tokens=2048,
            do_sample=True,
            temperature=0.7,
        )
        # Decode only the new tokens (skip the input)
        input_len = inputs["input_ids"].shape[1]
        response_tokens = outputs[0][input_len:]
        result: str = self._tokenizer.decode(response_tokens, skip_special_tokens=True)
        return result

    async def stream(self, prompt: str, system: str | None = None) -> AsyncIterator[str]:
        """Stream response from OpenVINO.

        Note: OpenVINO doesn't natively support token-by-token streaming
        with optimum-intel, so we generate the full response and yield it
        in chunks to maintain the streaming interface.
        """
        response = await self.generate(prompt, system)
        # Yield in word-sized chunks to simulate streaming
        words = response.split(" ")
        for i, word in enumerate(words):
            if i > 0:
                yield " "
            yield word
