"""Remote client for connecting to the NexaSci model server."""

from __future__ import annotations

from typing import Any, List, Optional, Sequence

import httpx

from .client_llm import Message


class RemoteNexaSciClient:
    """Client that connects to a remote NexaSci model server."""

    def __init__(self, base_url: str = "http://127.0.0.1:8001", timeout: float = 120.0) -> None:
        """Initialize the remote client.

        Parameters
        ----------
        base_url:
            Base URL of the model server (default: http://127.0.0.1:8001).
        timeout:
            Request timeout in seconds (default: 120.0 for long generations).
        """
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self._client = httpx.Client(timeout=timeout)

    def generate(
        self,
        messages: Sequence[Message],
        *,
        max_new_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
    ) -> str:
        """Generate a response from the remote model.

        Parameters
        ----------
        messages:
            Conversation history.
        max_new_tokens:
            Maximum tokens to generate.
        temperature:
            Sampling temperature.
        top_p:
            Top-p sampling parameter.

        Returns
        -------
        str
            Generated text response.
        """
        payload = {
            "messages": [{"role": msg.role, "content": msg.content} for msg in messages],
        }
        if max_new_tokens is not None:
            payload["max_new_tokens"] = max_new_tokens
        if temperature is not None:
            payload["temperature"] = temperature
        if top_p is not None:
            payload["top_p"] = top_p

        try:
            response = self._client.post(
                f"{self.base_url}/generate",
                json=payload,
            )
            response.raise_for_status()
            result = response.json()
            return result["text"]
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 503:
                raise RuntimeError("Model server not ready. Is the model loaded?") from e
            raise RuntimeError(f"Model server error: {e.response.text}") from e
        except httpx.RequestError as e:
            raise RuntimeError(f"Failed to connect to model server at {self.base_url}. Is it running?") from e

    @property
    def available_tools(self) -> Sequence[str]:
        """Return available tools from the model server."""
        try:
            response = self._client.get(f"{self.base_url}/tools")
            response.raise_for_status()
            result = response.json()
            return tuple(result.get("tools", []))
        except Exception:
            return tuple()

    def health_check(self) -> dict[str, Any]:
        """Check if the model server is healthy."""
        try:
            response = self._client.get(f"{self.base_url}/health")
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"status": "unhealthy", "error": str(e)}


__all__ = ["RemoteNexaSciClient"]