"""Pipeline integration for HuggingFace transformers.

Registers GrillyOptimum as a backend so HF pipelines can use it:
    from transformers import pipeline
    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
"""

from __future__ import annotations

import logging
from typing import Any

logger = logging.getLogger(__name__)


def create_text_generation_pipeline(
    model_id: str,
    tokenizer=None,
    dtype: str = "fp16",
    **kwargs,
):
    """Create a text-generation pipeline using Vulkan backend.

    Args:
        model_id: HF model ID or local path.
        tokenizer: Optional tokenizer (auto-loaded if None).
        dtype: Weight precision.
        **kwargs: Additional pipeline arguments.

    Returns:
        A callable pipeline object.
    """
    from grillyinference import TextGenerator, LlamaForCausalLM

    if tokenizer is None:
        from transformers import AutoTokenizer
        tokenizer = AutoTokenizer.from_pretrained(model_id)

    model = LlamaForCausalLM.from_pretrained(model_id, dtype=dtype)
    generator = TextGenerator(model, tokenizer)

    class VulkanTextGenerationPipeline:
        """Simple pipeline wrapper for text generation."""

        def __init__(self, generator, tokenizer):
            self._generator = generator
            self._tokenizer = tokenizer

        def __call__(
            self,
            text: str | list[str],
            max_new_tokens: int = 128,
            temperature: float = 0.7,
            top_k: int = 50,
            top_p: float = 0.9,
            **kwargs,
        ) -> list[dict[str, Any]]:
            if isinstance(text, str):
                text = [text]

            results = []
            for prompt in text:
                output = self._generator.generate(
                    prompt,
                    max_tokens=max_new_tokens,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                )
                results.append({
                    "generated_text": output,
                    "prompt": prompt,
                })
            return results

    return VulkanTextGenerationPipeline(generator, tokenizer)
