
# Wraps different model backends behind a unified ModelFn interface!!!! 

# Supported backends:
#   HuggingFace local (transformers pipeline or direct model)
#  OpenAI-compatible API (any local serverOllama, llama.cpp, LM Studio, vLLM)
#  Raw callable (pass-through for testing or custom wrappers)


import time
from typing import Optional
from .base import ModelFn


class ModelAdapter:
    """Base adapter. Subclass or use factory functions below."""

    def __init__(self, model_target: str):
        self.model_target = model_target

    def query(self, prompt: str, conversation_history: Optional[list[dict]] = None) -> str:
        """
        Query the model. Override in subclasses.
        conversation_history: list of {"role": "user"|"assistant", "content": str}
        Used by multi-turn attacks (crescendo, payload splitting).
        """
        raise NotImplementedError

    def is_reachable(self) -> tuple[bool, str]:
        """
        Preflight check. Returns (ok, reason).
        Called by runner before any attacks start.
        """
        try:
            response = self.query("Hello.")
            if response and not response.startswith("__ERROR__"):
                return True, "ok"
            return False, "Model returned empty response"
        except Exception as e:
            return False, str(e)

    def as_model_fn(self) -> ModelFn:
        """Return a simple ModelFn callable for single-turn attacks."""
        return lambda prompt: self.query(prompt)


class HuggingFaceAdapter(ModelAdapter):
    """
    Wraps a HuggingFace transformers pipeline or model+tokenizer pair.

    Usage:
        from transformers import pipeline
        pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
        adapter = HuggingFaceAdapter(pipe, model_target="TinyLlama-1.1B")
    """

    def __init__(self, pipeline_or_fn, model_target: str, max_new_tokens: int = 512):
        super().__init__(model_target)
        self._pipeline = pipeline_or_fn
        self._max_new_tokens = max_new_tokens

    def query(self, prompt: str, conversation_history: Optional[list[dict]] = None) -> str:
        try:
            # If conversation_history provided, format as chat template if possible
            if conversation_history:
                messages = conversation_history + [{"role": "user", "content": prompt}]
                try:
                    # Try chat template (works with instruct models)
                    outputs = self._pipeline(messages, max_new_tokens=self._max_new_tokens,
                                            return_full_text=False, do_sample=False)
                    return outputs[0]["generated_text"]
                except Exception:
                    # Fall back to plain prompt
                    pass

            outputs = self._pipeline(prompt, max_new_tokens=self._max_new_tokens,
                                     return_full_text=False, do_sample=False)
            if isinstance(outputs, list) and outputs:
                result = outputs[0]
                if isinstance(result, dict):
                    return result.get("generated_text", str(result))
                return str(result)
            return str(outputs)
        except Exception as e:
            return f"__ERROR__: {e}"


class OpenAICompatAdapter(ModelAdapter):
    """
    Wraps any OpenAI-compatible local API.
    Works with: Ollama, llama.cpp server, LM Studio, vLLM, etc.

    Usage:
        adapter = OpenAICompatAdapter(
            base_url="http://localhost:11434/v1",
            model_name="llama3",
            model_target="ollama/llama3",
        )
    """

    def __init__(
        self,
        base_url: str,
        model_name: str,
        model_target: str,
        api_key: str = "no-key",
        timeout: int = 60,
        max_tokens: int = 512,
        temperature: float = 0.0,
    ):
        super().__init__(model_target)
        self._base_url  = base_url.rstrip("/")
        self._model     = model_name
        self._api_key   = api_key
        self._timeout   = timeout
        self._max_tokens = max_tokens
        self._temperature = temperature

    def query(self, prompt: str, conversation_history: Optional[list[dict]] = None) -> str:
        try:
            import urllib.request, json as _json

            messages = (conversation_history or []) + [{"role": "user", "content": prompt}]
            payload = _json.dumps({
                "model":       self._model,
                "messages":    messages,
                "max_tokens":  self._max_tokens,
                "temperature": self._temperature,
            }).encode()

            req = urllib.request.Request(
                f"{self._base_url}/chat/completions",
                data=payload,
                headers={
                    "Content-Type":  "application/json",
                    "Authorization": f"Bearer {self._api_key}",
                },
            )
            with urllib.request.urlopen(req, timeout=self._timeout) as resp:
                data = _json.loads(resp.read())
                return data["choices"][0]["message"]["content"]
        except Exception as e:
            return f"__ERROR__: {e}"

    def is_reachable(self) -> tuple[bool, str]:
        try:
            import urllib.request
            req = urllib.request.Request(
                f"{self._base_url}/models",
                headers={"Authorization": f"Bearer {self._api_key}"},
            )
            with urllib.request.urlopen(req, timeout=5):
                pass
            return True, "ok"
        except Exception as e:
            return False, f"API not reachable at {self._base_url}: {e}"


class RawCallableAdapter(ModelAdapter):
    """
    Pass-through adapter for raw callables.
    Use when you already have a model_fn and just want preflight checking.
    """

    def __init__(self, fn: ModelFn, model_target: str):
        super().__init__(model_target)
        self._fn = fn

    def query(self, prompt: str, conversation_history: Optional[list[dict]] = None) -> str:
        try:
            return self._fn(prompt)
        except Exception as e:
            return f"__ERROR__: {e}"


class MLXAttackAdapter(ModelAdapter):
    """
    Attack-layer adapter backed by mlx-lm (Apple Silicon only).

    Drop-in replacement for HuggingFaceAdapter for text-generation models.
    Supports multi-turn attacks (crescendo, payload splitting) via chat templates.

    Usage:
        adapter = MLXAttackAdapter("Qwen/Qwen2.5-1.5B-Instruct", model_target="Qwen2.5-1.5B")
        response = adapter.query("Hello")
    """

    def __init__(self, model_id: str, model_target: str, max_tokens: int = 256):
        super().__init__(model_target)
        try:
            from sichgate.models.mlx_adapter import MLXModelAdapter
        except ImportError:
            raise RuntimeError("sichgate.models.mlx_adapter could not be imported")
        self._adapter = MLXModelAdapter(model_id=model_id, max_tokens=max_tokens)

    def query(self, prompt: str, conversation_history: Optional[list[dict]] = None) -> str:
        try:
            return self._adapter.query(prompt, conversation_history=conversation_history)
        except Exception as e:
            return f"__ERROR__: {e}"

    def is_reachable(self) -> tuple[bool, str]:
        return self._adapter.is_reachable()


# ── Auto-detect factory ───────────────────────────────────────────────────────

def make_adapter(model_target: str, **kwargs) -> ModelAdapter:
    """
    Auto-detect backend from model_target string and kwargs.

    Examples:
        make_adapter("TinyLlama/TinyLlama-1.1B-Chat-v1.0")          # HuggingFace
        make_adapter("./local/model/path")                            # HuggingFace local
        make_adapter("ollama/llama3", base_url="http://localhost:11434/v1")  # OpenAI-compat
        make_adapter("fn", fn=my_callable)                            # raw callable
    """
    if "fn" in kwargs:
        return RawCallableAdapter(fn=kwargs["fn"], model_target=model_target)

    if "base_url" in kwargs:
        model_name = kwargs.pop("model_name", model_target.split("/")[-1])
        return OpenAICompatAdapter(model_name=model_name, model_target=model_target, **kwargs)

    # Default: HuggingFace
    try:
        from transformers import pipeline as hf_pipeline
        pipe = hf_pipeline(
            "text-generation",
            model=model_target,
            max_new_tokens=kwargs.get("max_new_tokens", 512),
            do_sample=False,
        )
        return HuggingFaceAdapter(pipe, model_target=model_target,
                                  max_new_tokens=kwargs.get("max_new_tokens", 512))
    except Exception as e:
        raise RuntimeError(
            f"Could not load model '{model_target}'. "
            f"If using an API, pass base_url=... to make_adapter(). "
            f"Original error: {e}"
        )
