#  Universal model interface.

# Every attack talks to models through this adapter only.

from __future__ import annotations

import logging
from typing import Any, Optional

logger = logging.getLogger(__name__)


class ModelUnreachableError(RuntimeError):
   
    pass


class ModelAdapter:

    def __init__(
        self,
        model_path: str,
        adapter_type: str = "hf",
        *,
        device: Optional[str] = None,
        max_new_tokens: int = 256,
        temperature: float = 0.7,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        callable_fn: Optional[Any] = None,
        hf_token: Optional[str] = None,           # ADD: for private HF repos
        trust_remote_code: bool = False,           # ADD: for custom architectures
) -> None:
        self.model_path = model_path
        self.adapter_type = adapter_type
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self._pipeline = None
        self._callable = callable_fn
        self._base_url = base_url
        self._api_key = api_key
        self._device = device or self._detect_device()
        self._hf_token = hf_token or __import__('os').environ.get("HF_TOKEN")  
        self._trust_remote_code = trust_remote_code                            
 
    def query(self, prompt: str, **kwargs) -> str:
        try:
            if self.adapter_type == "callable":
                return self._query_callable(prompt, **kwargs)
            if self.adapter_type == "hf":
                return self._query_hf(prompt, **kwargs)
            if self.adapter_type == "openai_compat":
                return self._query_openai_compat(prompt, **kwargs)
            raise ModelUnreachableError(f"Unknown adapter_type: {self.adapter_type}")
        except ModelUnreachableError:
            raise
        except Exception as exc:
            raise ModelUnreachableError(f"Model query failed: {exc}") from exc

    def is_reachable(self) -> bool:
        try:
            self.query("Hello")
            return True
        except ModelUnreachableError:
            return False
        except Exception:
            return False

    def batch_query(self, prompts: list[str], **kwargs) -> list[str]:
        """Query multiple prompts. Falls back to serial if batch not supported."""
        return [self.query(p, **kwargs) for p in prompts]

    def _query_hf(self, prompt: str, **kwargs) -> str:
        if self._pipeline is None:
            quant_level = getattr(self, "_quantization_level", None)
            if quant_level is not None:
                self._pipeline = self._load_hf_pipeline_quantized(quant_level)
            else:
                self._pipeline = self._load_hf_pipeline()
        max_new_tokens = kwargs.get("max_new_tokens", self.max_new_tokens)
        output = self._pipeline(prompt, max_new_tokens=max_new_tokens, do_sample=True)
        if isinstance(output, list) and output:
            first = output[0]
            if isinstance(first, dict):
                generated = first.get("generated_text", "")
                # HuggingFace text-generation pipelines echo the prompt — strip it
                if generated.startswith(prompt):
                    generated = generated[len(prompt):]
                return generated.strip()
        return str(output)

    def _load_hf_pipeline(self):
        try:
            from transformers import pipeline as hf_pipeline
            import torch
        except ImportError as exc:
            raise ModelUnreachableError(
                "transformers / torch not installed. "
                "Run: pip install transformers torch"
            ) from exc

        device_map = self._device
        logger.info("Loading model %s on %s", self.model_path, device_map)
        try:
            pipe = hf_pipeline(
                "text-generation",
                model=self.model_path,
                device_map=device_map,
                torch_dtype="auto",
                token=self._hf_token,                      
                trust_remote_code=self._trust_remote_code, 

            )
            return pipe
        except Exception as exc:
            raise ModelUnreachableError(
                f"Could not load model at '{self.model_path}': {exc}"
            ) from exc

    def _query_callable(self, prompt: str, **kwargs) -> str:
        if self._callable is None:
            raise ModelUnreachableError("No callable provided to adapter")
        result = self._callable(prompt)
        return str(result)

    def _query_openai_compat(self, prompt: str, **kwargs) -> str:
        try:
            import httpx
        except ImportError as exc:
            raise ModelUnreachableError("httpx not installed") from exc

        if not self._base_url:
            raise ModelUnreachableError("base_url required for openai_compat adapter")

        headers = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"

        payload = {
            "model": self.model_path,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": kwargs.get("max_new_tokens", self.max_new_tokens),
            "temperature": kwargs.get("temperature", self.temperature),
        }
        response = httpx.post(
            f"{self._base_url}/chat/completions",
            json=payload,
            headers=headers,
            timeout=60,
        )
        response.raise_for_status()
        data = response.json()
        return data["choices"][0]["message"]["content"]

    @staticmethod
    def _detect_device() -> str:
        try:
            import torch
            if torch.cuda.is_available():
                return "cuda"
            if torch.backends.mps.is_available():
                return "mps"
        except ImportError:
            pass
        return "cpu"

    @classmethod
    def load_quantized(
        cls,
        model_path: str,
        quantization_level: str,
        *,
        hf_token: Optional[str] = None,
        trust_remote_code: bool = False,
        device: Optional[str] = None,
    ) -> "ModelAdapter":
        VALID_LEVELS = {"Q4", "Q8", "INT8", "FP16"}
        if quantization_level not in VALID_LEVELS:
            raise ValueError(
                f"Unknown quantization level '{quantization_level}'. "
                f"Valid options: {sorted(VALID_LEVELS)}"
            )

        adapter = cls(
            model_path=model_path,
            adapter_type="hf",
            device=device,
            hf_token=hf_token,
            trust_remote_code=trust_remote_code,
        )
        # Store the requested level so _load_hf_pipeline can apply it
        adapter._quantization_level = quantization_level
        return adapter

    def _load_hf_pipeline_quantized(self, quantization_level: str):
        try:
            from transformers import pipeline as hf_pipeline, BitsAndBytesConfig
            import torch
        except ImportError as exc:
            raise ModelUnreachableError(
                "transformers / torch not installed. "
                "Run: pip install transformers torch"
            ) from exc

        logger.info(
            "Loading model %s at %s quantization", self.model_path, quantization_level
        )

        kwargs: dict = {
            "model": self.model_path,
            "token": self._hf_token,
            "trust_remote_code": self._trust_remote_code,
        }

        try:
            if quantization_level in ("Q4",):
                try:
                    import bitsandbytes  # noqa: F401
                    bnb_config = BitsAndBytesConfig(load_in_4bit=True)
                    kwargs["quantization_config"] = bnb_config
                    kwargs["device_map"] = "auto"
                except ImportError:
                    logger.warning(
                        "bitsandbytes not installed; falling back to FP16 for Q4 request"
                    )
                    kwargs["torch_dtype"] = torch.float16
                    kwargs["device_map"] = self._device

            elif quantization_level in ("Q8", "INT8"):
                try:
                    import bitsandbytes  # noqa: F401
                    bnb_config = BitsAndBytesConfig(load_in_8bit=True)
                    kwargs["quantization_config"] = bnb_config
                    kwargs["device_map"] = "auto"
                except ImportError:
                    logger.warning(
                        "bitsandbytes not installed; falling back to FP16 for %s request",
                        quantization_level,
                    )
                    kwargs["torch_dtype"] = torch.float16
                    kwargs["device_map"] = self._device

            elif quantization_level == "FP16":
                kwargs["torch_dtype"] = torch.float16
                kwargs["device_map"] = self._device

            return hf_pipeline("text-generation", **kwargs)

        except Exception as exc:
            raise ModelUnreachableError(
                f"Could not load '{self.model_path}' at {quantization_level}: {exc}"
            ) from exc
