from __future__ import annotations

import json
import os
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

from prompt_injection_framework.models.base import BaseModelAdapter, ModelRequest


class OllamaModelAdapter(BaseModelAdapter):
    provider = "ollama"

    def __init__(self, model_name: str, base_url: str | None = None) -> None:
        super().__init__(model_name=model_name)
        configured_base = base_url or os.getenv("OLLAMA_BASE_URL") or "http://127.0.0.1:11434"
        self.base_url = configured_base.rstrip("/")
        self.api_url = f"{self.base_url}/api/generate"

    def _generate_text(self, request: ModelRequest) -> tuple[str, dict[str, object]]:
        prompt_sections: list[str] = []
        if request.context:
            prompt_sections.append(f"[SYSTEM]\n{request.context}")
        if request.conversation_history:
            history_lines = []
            for turn in request.conversation_history:
                history_lines.append(f"[{turn.role.upper()}]\n{turn.content}")
            prompt_sections.append("\n\n".join(history_lines))

        user_parts = [request.prompt]
        if request.task_input:
            user_parts.append(request.task_input)
        user_content = "\n\n".join(part for part in user_parts if part)
        prompt_sections.append(f"[USER]\n{user_content}")

        payload = {
            "model": self.model_name,
            "prompt": "\n\n".join(section for section in prompt_sections if section),
            "stream": False,
            "options": {
                "temperature": request.temperature,
                "num_predict": request.max_tokens,
            },
        }
        if request.seed is not None:
            payload["options"]["seed"] = request.seed

        req = Request(
            self.api_url,
            data=json.dumps(payload).encode("utf-8"),
            headers={
                "Content-Type": "application/json",
                "Accept": "application/json",
                "User-Agent": "llm-prompt-injection-security-eval-framework/0.1",
            },
            method="POST",
        )

        try:
            with urlopen(req, timeout=120) as response:
                raw_response = json.loads(response.read().decode("utf-8"))
        except HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"Ollama request failed with status {exc.code}: {detail}") from exc
        except URLError as exc:
            raise RuntimeError(f"Ollama request failed: {exc.reason}") from exc

        text = str(raw_response.get("response", ""))
        if not text.strip():
            raise RuntimeError("Ollama returned an empty response.")
        return text, raw_response