"""VulkanModelForCausalLM — HuggingFace-compatible model interface.

Wraps GrillyInference's LlamaForCausalLM with HF GenerationMixin patterns:
- from_pretrained() loads weights and creates inference engine
- generate() is compatible with HF generation kwargs
- Can be used with transformers.pipeline("text-generation")
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

import numpy as np

from .configuration import VulkanConfig

logger = logging.getLogger(__name__)


class VulkanModelForCausalLM:
    """HuggingFace-compatible Vulkan model for causal language modeling.

    Compatible with:
        from grillyoptimum import VulkanModelForCausalLM
        model = VulkanModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")
        output = model.generate(input_ids, max_new_tokens=100)
    """

    def __init__(self, model, config, vulkan_config=None):
        """Initialize with a GrillyInference model.

        Args:
            model: grillyinference.LlamaForCausalLM instance.
            config: grillyinference.LlamaConfig instance.
            vulkan_config: Optional VulkanConfig.
        """
        self._model = model
        self.config = config
        self.vulkan_config = vulkan_config or VulkanConfig()
        self._kv_cache = None

    @classmethod
    def from_pretrained(
        cls,
        model_id_or_path: str,
        dtype: str = "fp16",
        vulkan_config: VulkanConfig | None = None,
        **kwargs,
    ) -> VulkanModelForCausalLM:
        """Load a pretrained model, HuggingFace style.

        Args:
            model_id_or_path: HF model ID or local path.
            dtype: Weight precision.
            vulkan_config: Optional backend config.
            **kwargs: Additional arguments (ignored for HF compat).

        Returns:
            VulkanModelForCausalLM instance.
        """
        from grillyinference import LlamaForCausalLM, LlamaConfig

        if vulkan_config is None:
            vulkan_config = VulkanConfig(dtype=dtype)

        config = LlamaConfig.from_pretrained(model_id_or_path)
        model = LlamaForCausalLM.from_pretrained(
            model_id_or_path,
            dtype=vulkan_config.dtype,
        )

        return cls(model, config, vulkan_config)

    def generate(
        self,
        input_ids: np.ndarray | None = None,
        max_new_tokens: int = 128,
        max_length: int | None = None,
        temperature: float = 1.0,
        top_k: int = 50,
        top_p: float = 1.0,
        do_sample: bool = True,
        num_return_sequences: int = 1,
        **kwargs,
    ) -> np.ndarray:
        """Generate tokens, HuggingFace style.

        Args:
            input_ids: (batch, seq_len) int32 input token IDs.
            max_new_tokens: Maximum tokens to generate.
            max_length: Maximum total length (alternative to max_new_tokens).
            temperature: Sampling temperature.
            top_k: Top-k filtering.
            top_p: Nucleus sampling.
            do_sample: Whether to sample (True) or greedy decode (False).
            **kwargs: Additional HF generation kwargs (ignored).

        Returns:
            (batch, seq_len + generated) int32 array of token IDs.
        """
        if input_ids is None:
            raise ValueError("input_ids is required")

        if isinstance(input_ids, list):
            input_ids = np.array(input_ids, dtype=np.int32)
        if input_ids.ndim == 1:
            input_ids = input_ids[np.newaxis, :]

        batch_size = input_ids.shape[0]
        if batch_size != 1:
            raise ValueError("Only batch_size=1 supported currently")

        if max_length is not None:
            max_new_tokens = max_length - input_ids.shape[1]

        if not do_sample:
            temperature = 0.0

        # Import sampling function
        from grillyinference.inference.generate import _sample_top_k_top_p, LLAMA3_STOP_TOKENS

        # Prefill
        from grillyinference import KVCache
        kv_cache = KVCache(
            self.config,
            max_batch=1,
            page_size=self.vulkan_config.page_size,
            raw_window=self.vulkan_config.raw_window,
            h2o_lambda=self.vulkan_config.h2o_lambda,
            enable_vsa=self.vulkan_config.enable_vsa,
        )

        logits = self._model.forward(input_ids, kv_cache=kv_cache)
        last_logits = logits[0, -1, :]

        generated = list(input_ids[0])

        for _ in range(max_new_tokens):
            token_id = _sample_top_k_top_p(last_logits, temperature, top_k, top_p)
            if token_id in LLAMA3_STOP_TOKENS:
                break
            generated.append(token_id)
            token_array = np.array([[token_id]], dtype=np.int32)
            logits = self._model.decode_step(token_array, kv_cache)
            last_logits = logits[0, -1, :]

        return np.array([generated], dtype=np.int32)

    def __call__(self, input_ids, **kwargs):
        """Make model callable for pipeline compatibility."""
        return self.generate(input_ids, **kwargs)

    @property
    def device(self):
        return self.vulkan_config.device

    def memory_footprint(self) -> dict:
        return self._model.memory_footprint()
