"""LLM provider for custom language models like Pacific-Prime."""
import asyncio
from typing import Optional
from .base import BaseProvider, OracleData


class LLMProvider(BaseProvider):
    """LLM oracle provider for text generation.

    Supports HuggingFace models including custom architectures
    like Pacific-Prime's INL architecture.
    """

    name = "llm"

    def __init__(self, model_name: str = "Pacific-Prime/pacific-prime", max_tokens: int = 100):
        self.model_name = model_name
        self.max_tokens = max_tokens
        self._model = None
        self._tokenizer = None

    def _load_model(self):
        """Lazy load the LLM model."""
        if self._model is not None:
            return

        try:
            from transformers import AutoModelForCausalLM, AutoTokenizer
            import torch

            print(f"[LLM] Loading model: {self.model_name}")
            print("[LLM] This may take a moment...")

            # Determine device and dtype
            if torch.cuda.is_available():
                device = "cuda"
                dtype = torch.bfloat16
            else:
                device = "cpu"
                dtype = torch.float32

            self._tokenizer = AutoTokenizer.from_pretrained(
                self.model_name,
                trust_remote_code=True,
            )

            self._model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                trust_remote_code=True,
                torch_dtype=dtype,
            ).to(device)

            # Set pad token if not set
            if self._tokenizer.pad_token is None:
                self._tokenizer.pad_token = self._tokenizer.eos_token

            print(f"[LLM] Model loaded on {device}")

        except ImportError:
            raise ImportError("transformers and torch required. Install with: pip install nanopy-oracle[ai]")

    def _generate(self, prompt: str) -> str:
        """Generate text from prompt."""
        import torch

        device = next(self._model.parameters()).device
        inputs = self._tokenizer(prompt, return_tensors="pt", padding=True).to(device)
        input_length = inputs["input_ids"].shape[1]

        with torch.no_grad():
            outputs = self._model.generate(
                **inputs,
                max_new_tokens=self.max_tokens,
                do_sample=True,
                temperature=0.8,
                top_p=0.92,
                top_k=50,
                pad_token_id=self._tokenizer.pad_token_id,
                eos_token_id=self._tokenizer.eos_token_id,
                use_cache=True,
            )

        # Decode full output then remove prompt
        full_text = self._tokenizer.decode(outputs[0], skip_special_tokens=True)

        # Remove the prompt from the output
        prompt_decoded = self._tokenizer.decode(inputs["input_ids"][0], skip_special_tokens=True)
        if full_text.startswith(prompt_decoded):
            text = full_text[len(prompt_decoded):]
        else:
            # Fallback: just decode new tokens
            generated = outputs[0][input_length:]
            text = self._tokenizer.decode(generated, skip_special_tokens=True)

        return text.strip()

    async def fetch(self, key: str, prompt: Optional[str] = None, **kwargs) -> OracleData:
        """Generate text using the LLM.

        Args:
            key: Identifier for this generation (e.g., "chat:nanopy")
            prompt: The prompt to generate from
            **kwargs: Additional arguments

        Returns:
            OracleData with generated text
        """
        if prompt is None:
            prompt = kwargs.get("default_prompt", "Hello, I am an AI oracle.")

        self._load_model()

        # Run generation in thread pool
        loop = asyncio.get_event_loop()
        generated_text = await loop.run_in_executor(None, self._generate, prompt)

        # Truncate for on-chain storage (keep reasonable size)
        max_chars = 500
        if len(generated_text) > max_chars:
            generated_text = generated_text[:max_chars] + "..."

        return self.create_data(
            key=key,
            value=generated_text,
            confidence=0.9,  # LLM confidence is subjective
            metadata={
                "model": self.model_name,
                "prompt_length": len(prompt),
                "response_length": len(generated_text),
                "max_tokens": self.max_tokens,
            },
        )

    async def validate(self, data: OracleData) -> bool:
        """Validate LLM output."""
        if not isinstance(data.value, str):
            return False
        if len(data.value) == 0:
            return False
        return True


class ChatProvider(LLMProvider):
    """Chat-style LLM provider with conversation format."""

    name = "chat"

    def _format_chat(self, user_message: str, system_prompt: Optional[str] = None) -> str:
        """Format message for chat models."""
        messages = []

        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})

        messages.append({"role": "user", "content": user_message})

        # Try to use chat template if available
        if hasattr(self._tokenizer, "apply_chat_template"):
            try:
                return self._tokenizer.apply_chat_template(
                    messages,
                    tokenize=False,
                    add_generation_prompt=True,
                )
            except Exception:
                pass

        # Fallback for Pacific-Prime format
        formatted = ""
        if system_prompt:
            formatted += f"<SYSTEM>{system_prompt}</SYSTEM>\n"
        formatted += f"<USER>{user_message}</USER>\n<ASSISTANT>"
        return formatted

    async def fetch(self, key: str, prompt: Optional[str] = None, **kwargs) -> OracleData:
        """Generate chat response.

        Args:
            key: Identifier for this chat (e.g., "chat:question")
            prompt: User message
            **kwargs: system_prompt for system instructions

        Returns:
            OracleData with assistant response
        """
        if prompt is None:
            prompt = "What is NanoPy?"

        system_prompt = kwargs.get("system_prompt", "You are a helpful AI assistant.")

        self._load_model()

        # Format as chat
        formatted_prompt = self._format_chat(prompt, system_prompt)

        # Generate
        loop = asyncio.get_event_loop()
        generated_text = await loop.run_in_executor(None, self._generate, formatted_prompt)

        # Clean up response (remove trailing tokens)
        for stop_token in ["</ASSISTANT>", "<USER>", "<SYSTEM>", "<ERROR>"]:
            if stop_token in generated_text:
                generated_text = generated_text.split(stop_token)[0]

        generated_text = generated_text.strip()

        # Truncate for on-chain
        max_chars = 500
        if len(generated_text) > max_chars:
            generated_text = generated_text[:max_chars] + "..."

        return self.create_data(
            key=key,
            value=generated_text,
            confidence=0.9,
            metadata={
                "model": self.model_name,
                "prompt": prompt[:100],  # Store truncated prompt
                "response_length": len(generated_text),
                "type": "chat",
            },
        )
