"""
LM Studio Adapter — Talk to any OpenAI-compatible local LLM server from Python.

Works with: LM Studio, LocalAI, text-generation-webui (openai extension),
llama.cpp server, Ollama (OpenAI mode), vLLM, FastChat, TabbyAPI, Aphrodite,
or anything that serves /v1/chat/completions and /v1/completions.

Uses ONLY Python stdlib (urllib, json). No openai package required.

Usage:
    from brynq.lmstudio_adapter import is_server_running, list_models, chat, complete

    if is_server_running():
        models = list_models()
        reply = chat(models[0], [{"role": "user", "content": "Hello!"}])
        print(reply)
"""

import json
import urllib.error
import urllib.request
from typing import Optional

# Default timeout for all requests (seconds).
_TIMEOUT = 120


# ── Errors ────────────────────────────────────────────────────────────────────


class LMStudioConnectionError(ConnectionError):
    """Raised when the local LLM server is unreachable."""


class LMStudioAPIError(RuntimeError):
    """Raised when the server returns a non-200 HTTP status."""


class LMStudioModelNotFoundError(LMStudioAPIError):
    """Raised when the requested model is not loaded on the server."""


# ── Internal helpers ──────────────────────────────────────────────────────────


def _base_url(port: int) -> str:
    """Build the base URL for the local server."""
    return f"http://localhost:{port}/v1"


def _http_get(url: str, timeout: int = _TIMEOUT) -> dict:
    """GET *url* and return the parsed JSON response.

    Raises:
        LMStudioConnectionError: Server is not reachable.
        LMStudioAPIError: Server returned an HTTP error.
    """
    req = urllib.request.Request(
        url,
        headers={"Accept": "application/json"},
        method="GET",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.URLError as exc:
        raise LMStudioConnectionError(
            f"Cannot connect to local LLM server at {url}. "
            "Is LM Studio (or another OpenAI-compatible server) running? "
            "Start it and ensure the local server is enabled."
        ) from exc
    except urllib.error.HTTPError as exc:
        detail = ""
        try:
            detail = exc.read().decode("utf-8", errors="replace")
        except Exception:
            pass
        raise LMStudioAPIError(
            f"GET {url} returned HTTP {exc.code}: {detail}"
        ) from exc


def _http_post(url: str, body: dict, timeout: int = _TIMEOUT) -> dict:
    """POST JSON to *url* and return the parsed JSON response.

    Raises:
        LMStudioConnectionError: Server is not reachable.
        LMStudioAPIError: Server returned an HTTP error.
        LMStudioModelNotFoundError: The requested model is not available.
    """
    data = json.dumps(body).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=data,
        headers={
            "Content-Type": "application/json",
            "Accept": "application/json",
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = ""
        try:
            detail = exc.read().decode("utf-8", errors="replace")
        except Exception:
            pass
        # Detect model-not-found (404 or message containing "model")
        if exc.code == 404 or "model" in detail.lower():
            available = _try_list_models(url)
            hint = (
                f" Available models: {available}" if available
                else " Could not retrieve available models."
            )
            raise LMStudioModelNotFoundError(
                f"Model not found.{hint} (HTTP {exc.code}: {detail})"
            ) from exc
        raise LMStudioAPIError(
            f"POST {url} returned HTTP {exc.code}: {detail}"
        ) from exc
    except urllib.error.URLError as exc:
        raise LMStudioConnectionError(
            f"Cannot connect to local LLM server at {url}. "
            "Is LM Studio (or another OpenAI-compatible server) running? "
            "Start it and ensure the local server is enabled."
        ) from exc


def _try_list_models(url: str) -> list[str]:
    """Best-effort attempt to list models, used for error hints."""
    try:
        # Derive the models endpoint from any /v1/* URL
        base = url.rsplit("/v1/", 1)[0] + "/v1"
        data = _http_get(f"{base}/models", timeout=5)
        return [m.get("id", "?") for m in data.get("data", [])]
    except Exception:
        return []


# ── Public API ────────────────────────────────────────────────────────────────


def is_server_running(port: int = 1234) -> bool:
    """Check whether a local LLM server is reachable.

    Sends GET /v1/models and returns True if the server responds with HTTP 200.

    Args:
        port: Server port (default 1234, the LM Studio default).

    Returns:
        True if the server is running and responding, False otherwise.
    """
    try:
        _http_get(f"{_base_url(port)}/models", timeout=5)
        return True
    except (LMStudioConnectionError, LMStudioAPIError):
        return False


def list_models(port: int = 1234) -> list[str]:
    """List all models available on the local server.

    Calls GET /v1/models and returns a list of model ID strings.

    Args:
        port: Server port (default 1234).

    Returns:
        List of model identifier strings (e.g. ["lmstudio-community/Meta-Llama-3-8B-Instruct-GGUF"]).

    Raises:
        LMStudioConnectionError: Server is not reachable.
        LMStudioAPIError: Server returned an error.
    """
    data = _http_get(f"{_base_url(port)}/models")
    return [m.get("id", "") for m in data.get("data", []) if m.get("id")]


def chat(
    model: str,
    messages: list[dict],
    port: int = 1234,
    temperature: float = 0.7,
    max_tokens: int = 2048,
) -> str:
    """Send a chat completion request and return the assistant's reply.

    Uses the OpenAI-compatible POST /v1/chat/completions endpoint.

    Args:
        model: Model identifier (from list_models()).
        messages: Conversation in OpenAI format, e.g.
            [{"role": "user", "content": "Hello"}].
        port: Server port (default 1234).
        temperature: Sampling temperature (0.0 = deterministic, higher = more creative).
        max_tokens: Maximum tokens to generate.

    Returns:
        The assistant message content as a string.

    Raises:
        LMStudioConnectionError: Server is not reachable.
        LMStudioModelNotFoundError: The requested model is not loaded.
        LMStudioAPIError: Server returned an error.
        ValueError: Response is missing expected fields.
    """
    url = f"{_base_url(port)}/chat/completions"
    body = {
        "model": model,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    data = _http_post(url, body)

    # Extract the assistant content from the response
    choices = data.get("choices", [])
    if not choices:
        raise ValueError(
            f"Server returned no choices. Full response: {json.dumps(data)[:500]}"
        )
    content = choices[0].get("message", {}).get("content")
    if content is None:
        raise ValueError(
            f"No content in first choice. Full response: {json.dumps(data)[:500]}"
        )
    return content


def complete(
    model: str,
    prompt: str,
    port: int = 1234,
    temperature: float = 0.7,
    max_tokens: int = 2048,
) -> str:
    """Send a text completion request and return the generated text.

    Uses the OpenAI-compatible POST /v1/completions endpoint.
    This is a simpler interface for single-turn prompts without
    chat roles.

    Args:
        model: Model identifier (from list_models()).
        prompt: The text prompt to complete.
        port: Server port (default 1234).
        temperature: Sampling temperature.
        max_tokens: Maximum tokens to generate.

    Returns:
        The generated text as a string.

    Raises:
        LMStudioConnectionError: Server is not reachable.
        LMStudioModelNotFoundError: The requested model is not loaded.
        LMStudioAPIError: Server returned an error.
        ValueError: Response is missing expected fields.
    """
    url = f"{_base_url(port)}/completions"
    body = {
        "model": model,
        "prompt": prompt,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    data = _http_post(url, body)

    choices = data.get("choices", [])
    if not choices:
        raise ValueError(
            f"Server returned no choices. Full response: {json.dumps(data)[:500]}"
        )
    text = choices[0].get("text")
    if text is None:
        raise ValueError(
            f"No text in first choice. Full response: {json.dumps(data)[:500]}"
        )
    return text
