"""Unified inference client for sageLLM REST API.

Provides a synchronous Python client for interacting with a running
sageLLM server (``sagellm serve``).

Example::

    from sagellm import UnifiedInferenceClient

    client = UnifiedInferenceClient()  # connects to SagePorts.SAGELLM_SERVE_PORT
    response = client.complete(prompt="Hello!", max_tokens=128)
    print(response["text"])
"""

from __future__ import annotations

import json
import logging
from typing import Any
from urllib import error as urllib_error
from urllib import request as urllib_request

from sage.common.config.ports import SagePorts

logger = logging.getLogger(__name__)

_DEFAULT_LLM_URL: str = f"http://localhost:{SagePorts.SAGELLM_SERVE_PORT}"


class UnifiedInferenceClient:
    """Synchronous client for the sageLLM OpenAI-compatible REST API.

    Wraps ``/v1/chat/completions`` (and more) so that users can interact
    with a running ``sagellm serve`` instance from plain Python without
    installing additional HTTP libraries.

    Args:
        base_url: Base URL of the sageLLM server
            (default: ``http://localhost:{SagePorts.SAGELLM_SERVE_PORT}``).
        model: Default model name. Can be overridden per-request.
        api_key: Optional API key (sent as ``Authorization: Bearer <key>``).
        timeout_s: Default request timeout in seconds.
        offline_mode: When ``True``, return canned responses instead of
            making real HTTP calls.  Useful for demos and tests.

    Example::

        client = UnifiedInferenceClient()  # default: SagePorts.SAGELLM_SERVE_PORT
        resp = client.complete("Tell me a joke", max_tokens=50)
        print(resp["text"])
    """

    def __init__(
        self,
        base_url: str = _DEFAULT_LLM_URL,
        *,
        model: str | None = None,
        api_key: str | None = None,
        timeout_s: float = 60.0,
        offline_mode: bool = False,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._model = model
        self._api_key = api_key
        self._timeout_s = timeout_s
        self._offline_mode = offline_mode

    # ------------------------------------------------------------------
    # Public helpers
    # ------------------------------------------------------------------

    @property
    def base_url(self) -> str:
        """Return the configured base URL."""
        return self._base_url

    # ------------------------------------------------------------------
    # Core API
    # ------------------------------------------------------------------

    def complete(
        self,
        prompt: str,
        *,
        model: str | None = None,
        max_tokens: int = 128,
        temperature: float = 0.0,
        top_p: float = 1.0,
        stream: bool = False,
    ) -> dict[str, Any]:
        """Send a completion request to the server.

        This wraps the ``/v1/chat/completions`` endpoint, converting a
        simple prompt string into a ``[{"role": "user", "content": prompt}]``
        message list.

        Args:
            prompt: User prompt text.
            model: Model override (falls back to ``self._model``).
            max_tokens: Maximum tokens to generate.
            temperature: Sampling temperature.
            top_p: Nucleus-sampling probability mass.
            stream: Whether to request streaming (not yet supported).

        Returns:
            A dict with at least ``"text"`` (the generated text) and
            ``"usage"`` (token counts).
        """
        if self._offline_mode:
            return self._offline_complete(prompt, max_tokens)

        resolved_model = model or self._model or "default"

        payload: dict[str, Any] = {
            "model": resolved_model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "stream": False,  # streaming handled separately later
        }

        data = self._post("/v1/chat/completions", payload)

        # Extract text from OpenAI-compatible response
        choices = data.get("choices", [])
        text = choices[0]["message"]["content"] if choices else ""
        usage = data.get("usage", {})

        return {
            "text": text,
            "usage": usage,
            "raw": data,
        }

    def chat(
        self,
        messages: list[dict[str, str]],
        *,
        model: str | None = None,
        max_tokens: int = 128,
        temperature: float = 0.0,
        top_p: float = 1.0,
    ) -> dict[str, Any]:
        """Send a multi-turn chat request.

        Args:
            messages: List of ``{"role": ..., "content": ...}`` dicts.
            model: Model override.
            max_tokens: Maximum tokens to generate.
            temperature: Sampling temperature.
            top_p: Nucleus-sampling probability mass.

        Returns:
            A dict with ``"text"``, ``"usage"``, and ``"raw"`` keys.
        """
        if self._offline_mode:
            prompt = messages[-1].get("content", "") if messages else ""
            return self._offline_complete(prompt, max_tokens)

        resolved_model = model or self._model or "default"

        payload: dict[str, Any] = {
            "model": resolved_model,
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "stream": False,
        }

        data = self._post("/v1/chat/completions", payload)
        choices = data.get("choices", [])
        text = choices[0]["message"]["content"] if choices else ""
        usage = data.get("usage", {})

        return {
            "text": text,
            "usage": usage,
            "raw": data,
        }

    def health(self) -> bool:
        """Check server health.

        Returns:
            ``True`` if the server is reachable and healthy.
        """
        if self._offline_mode:
            return True
        try:
            self._get("/health")
            return True
        except Exception:
            return False

    # ------------------------------------------------------------------
    # HTTP helpers
    # ------------------------------------------------------------------

    def _headers(self) -> dict[str, str]:
        headers: dict[str, str] = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
        return headers

    def _post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
        url = f"{self._base_url}{path}"
        req = urllib_request.Request(
            url=url,
            data=json.dumps(payload).encode("utf-8"),
            headers=self._headers(),
            method="POST",
        )
        try:
            with urllib_request.urlopen(req, timeout=self._timeout_s) as resp:  # noqa: S310
                return json.loads(resp.read().decode("utf-8"))
        except urllib_error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"sageLLM server returned HTTP {exc.code}: {body}") from exc
        except urllib_error.URLError as exc:
            raise ConnectionError(
                f"Cannot connect to sageLLM server at {self._base_url}: {exc.reason}"
            ) from exc

    def _get(self, path: str) -> dict[str, Any]:
        url = f"{self._base_url}{path}"
        req = urllib_request.Request(url=url, headers=self._headers(), method="GET")
        try:
            with urllib_request.urlopen(req, timeout=self._timeout_s) as resp:  # noqa: S310
                return json.loads(resp.read().decode("utf-8"))
        except urllib_error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"sageLLM server returned HTTP {exc.code}: {body}") from exc
        except urllib_error.URLError as exc:
            raise ConnectionError(
                f"Cannot connect to sageLLM server at {self._base_url}: {exc.reason}"
            ) from exc

    # ------------------------------------------------------------------
    # Offline / demo helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _offline_complete(prompt: str, max_tokens: int) -> dict[str, Any]:
        text = f"[offline] Echo: {prompt[:50]}..."
        return {
            "text": text,
            "usage": {
                "prompt_tokens": len(prompt.split()),
                "completion_tokens": min(max_tokens, 10),
                "total_tokens": len(prompt.split()) + min(max_tokens, 10),
            },
            "raw": {"offline": True},
        }

    def __repr__(self) -> str:
        mode = " offline=True" if self._offline_mode else ""
        return f"UnifiedInferenceClient({self._base_url!r}{mode})"
