"""
Ollama Adapter -- local LLM access for Brynq agents via Ollama's REST API.

Pure stdlib implementation (urllib, json, http.client). No external dependencies.

Provides five core functions:
    is_ollama_running()  -- health check
    list_models()        -- enumerate available models
    chat()               -- multi-turn chat completions (streaming + non-streaming)
    generate()           -- single-prompt generation (streaming + non-streaming)
    get_model_info()     -- model metadata and parameters

Usage:
    from brynq.ollama_adapter import is_ollama_running, chat, list_models

    if is_ollama_running():
        models = list_models()
        reply = chat("llama3:latest", [{"role": "user", "content": "Hello"}])
        print(reply)
"""

from __future__ import annotations

import http.client
import json
import socket
import urllib.error
import urllib.request
from typing import Any, Generator, Optional

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

OLLAMA_HOST = "localhost"
OLLAMA_PORT = 11434
OLLAMA_BASE = f"http://{OLLAMA_HOST}:{OLLAMA_PORT}"
DEFAULT_TIMEOUT = 120  # seconds -- large models can be slow on first load


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _url(path: str) -> str:
    """Build a full Ollama URL from a path fragment."""
    return f"{OLLAMA_BASE}{path}"


def _get_json(path: str, *, timeout: int = DEFAULT_TIMEOUT) -> Any:
    """HTTP GET that returns parsed JSON.

    Raises:
        ConnectionError: If Ollama is unreachable.
        RuntimeError: On non-200 HTTP responses.
    """
    url = _url(path)
    req = urllib.request.Request(url, method="GET", headers={"Accept": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8")
            return json.loads(raw) if raw.strip() else {}
    except (urllib.error.URLError, ConnectionRefusedError, socket.timeout, OSError) as exc:
        raise ConnectionError(
            f"Cannot reach Ollama at {OLLAMA_BASE}. "
            f"Is Ollama running? Start it with: ollama serve\n"
            f"Detail: {exc}"
        ) from exc


def _post_json(
    path: str,
    body: dict,
    *,
    timeout: int = DEFAULT_TIMEOUT,
) -> Any:
    """HTTP POST with a JSON body, returning parsed JSON.

    Raises:
        ConnectionError: If Ollama is unreachable.
        RuntimeError: On non-200 HTTP responses.
    """
    url = _url(path)
    data = json.dumps(body).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=data,
        method="POST",
        headers={"Content-Type": "application/json", "Accept": "application/json"},
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8")
            return json.loads(raw) if raw.strip() else {}
    except urllib.error.HTTPError as exc:
        detail = ""
        try:
            detail = exc.read().decode("utf-8", errors="replace")
        except Exception:
            pass
        raise RuntimeError(
            f"Ollama returned HTTP {exc.code} for POST {path}: {detail}"
        ) from exc
    except (urllib.error.URLError, ConnectionRefusedError, socket.timeout, OSError) as exc:
        raise ConnectionError(
            f"Cannot reach Ollama at {OLLAMA_BASE}. "
            f"Is Ollama running? Start it with: ollama serve\n"
            f"Detail: {exc}"
        ) from exc


def _post_stream(
    path: str,
    body: dict,
    *,
    timeout: int = DEFAULT_TIMEOUT,
) -> Generator[dict, None, None]:
    """HTTP POST that yields parsed JSON objects from a newline-delimited stream.

    Ollama streams responses as one JSON object per line. This generator
    yields each parsed object as it arrives, enabling real-time token display.

    Raises:
        ConnectionError: If Ollama is unreachable.
        RuntimeError: On non-200 HTTP responses.
    """
    # Use http.client for true streaming -- urllib.request buffers the
    # entire response before returning, which defeats the purpose.
    body["stream"] = True
    payload = json.dumps(body).encode("utf-8")

    try:
        conn = http.client.HTTPConnection(OLLAMA_HOST, OLLAMA_PORT, timeout=timeout)
        conn.request(
            "POST",
            path,
            body=payload,
            headers={"Content-Type": "application/json"},
        )
        resp = conn.getresponse()
    except (ConnectionRefusedError, socket.timeout, OSError) as exc:
        raise ConnectionError(
            f"Cannot reach Ollama at {OLLAMA_BASE}. "
            f"Is Ollama running? Start it with: ollama serve\n"
            f"Detail: {exc}"
        ) from exc

    if resp.status != 200:
        detail = resp.read().decode("utf-8", errors="replace")
        conn.close()
        raise RuntimeError(
            f"Ollama returned HTTP {resp.status} for POST {path}: {detail}"
        )

    try:
        buffer = b""
        while True:
            chunk = resp.read(4096)
            if not chunk:
                # Process any remaining data in buffer
                if buffer.strip():
                    try:
                        yield json.loads(buffer.decode("utf-8"))
                    except json.JSONDecodeError:
                        pass
                break

            buffer += chunk
            # Split on newlines -- each line is a complete JSON object
            while b"\n" in buffer:
                line, buffer = buffer.split(b"\n", 1)
                line = line.strip()
                if line:
                    try:
                        yield json.loads(line.decode("utf-8"))
                    except json.JSONDecodeError:
                        pass
    finally:
        conn.close()


def _ensure_ollama() -> None:
    """Raise ConnectionError if Ollama is not reachable."""
    if not is_ollama_running():
        raise ConnectionError(
            f"Ollama is not running at {OLLAMA_BASE}. "
            "Start it with: ollama serve"
        )


def _validate_model(model: str) -> None:
    """Raise ValueError if the model is not available, listing what is.

    Assumes Ollama is already confirmed reachable.
    """
    try:
        available = list_models()
    except Exception:
        return  # Cannot check -- let Ollama return its own error

    names = [m["name"] for m in available]
    if model not in names:
        # Also check without :latest tag
        bare_names = [n.split(":")[0] for n in names]
        if model.split(":")[0] not in bare_names:
            models_str = ", ".join(names) if names else "(none installed)"
            raise ValueError(
                f"Model '{model}' not found. "
                f"Available models: {models_str}. "
                f"Pull it with: ollama pull {model}"
            )


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def is_ollama_running() -> bool:
    """Check whether Ollama is running at localhost:11434.

    Sends a GET to http://localhost:11434/ and returns True if the
    server responds with HTTP 200.

    Returns:
        True if Ollama is reachable, False otherwise.

    Example:
        >>> if is_ollama_running():
        ...     print("Ollama is up")
    """
    try:
        req = urllib.request.Request(_url("/"), method="GET")
        with urllib.request.urlopen(req, timeout=5) as resp:
            return resp.status == 200
    except Exception:
        return False


def list_models() -> list[dict]:
    """List all models available in the local Ollama installation.

    Queries GET /api/tags and returns a simplified list of model metadata.

    Returns:
        List of dicts, each with at minimum:
            - name (str): Model name with tag, e.g. "llama3:latest"
            - size (int): Model size in bytes
            - modified_at (str): Last modification timestamp

    Raises:
        ConnectionError: If Ollama is not running.

    Example:
        >>> models = list_models()
        >>> for m in models:
        ...     print(f"{m['name']}  ({m['size'] / 1e9:.1f} GB)")
    """
    _ensure_ollama()
    data = _get_json("/api/tags")
    models = data.get("models", [])
    return [
        {
            "name": m.get("name", "unknown"),
            "size": m.get("size", 0),
            "modified_at": m.get("modified_at", ""),
            "digest": m.get("digest", ""),
            "details": m.get("details", {}),
        }
        for m in models
    ]


def chat(
    model: str,
    messages: list[dict],
    *,
    stream: bool = False,
    timeout: int = DEFAULT_TIMEOUT,
    **kwargs: Any,
) -> str | Generator[str, None, None]:
    """Send a chat completion request to Ollama.

    Multi-turn conversation interface. Messages use the standard format:
        [{"role": "system", "content": "..."}, {"role": "user", "content": "..."}]

    Args:
        model: Ollama model name, e.g. "llama3:latest" or "mistral".
        messages: Conversation history as a list of role/content dicts.
            Supported roles: "system", "user", "assistant".
        stream: If False (default), block until the full response is ready
            and return the complete text as a string. If True, return a
            generator that yields response text chunks as they arrive.
        timeout: Request timeout in seconds (default 120).
        **kwargs: Additional parameters forwarded to Ollama (e.g.
            temperature, top_p, num_predict). See Ollama API docs.

    Returns:
        If stream=False: the full response text (str).
        If stream=True: a generator yielding text chunks (str).

    Raises:
        ConnectionError: If Ollama is not running.
        ValueError: If the model is not found (includes list of available models).
        RuntimeError: On other Ollama API errors.

    Example (non-streaming):
        >>> reply = chat("llama3", [{"role": "user", "content": "What is 2+2?"}])
        >>> print(reply)

    Example (streaming):
        >>> for chunk in chat("llama3", [{"role": "user", "content": "Hi"}], stream=True):
        ...     print(chunk, end="", flush=True)
    """
    _ensure_ollama()
    _validate_model(model)

    body: dict[str, Any] = {
        "model": model,
        "messages": messages,
        "stream": stream,
    }
    body.update(kwargs)

    if stream:
        return _chat_stream(body, timeout=timeout)
    else:
        resp = _post_json("/api/chat", body, timeout=timeout)
        return resp.get("message", {}).get("content", "")


def _chat_stream(body: dict, *, timeout: int) -> Generator[str, None, None]:
    """Internal generator for streaming chat responses."""
    for obj in _post_stream("/api/chat", body, timeout=timeout):
        content = obj.get("message", {}).get("content", "")
        if content:
            yield content


def generate(
    model: str,
    prompt: str,
    *,
    stream: bool = False,
    timeout: int = DEFAULT_TIMEOUT,
    **kwargs: Any,
) -> str | Generator[str, None, None]:
    """Generate a completion from a single prompt.

    Simpler interface than chat() for one-shot generation without
    conversation history.

    Args:
        model: Ollama model name, e.g. "llama3:latest" or "codellama".
        prompt: The input prompt text.
        stream: If False (default), block and return the full text.
            If True, return a generator yielding text chunks.
        timeout: Request timeout in seconds (default 120).
        **kwargs: Additional Ollama parameters (temperature, top_p, etc.).

    Returns:
        If stream=False: the full response text (str).
        If stream=True: a generator yielding text chunks (str).

    Raises:
        ConnectionError: If Ollama is not running.
        ValueError: If the model is not found.
        RuntimeError: On other Ollama API errors.

    Example:
        >>> result = generate("llama3", "Explain Python generators in one paragraph.")
        >>> print(result)
    """
    _ensure_ollama()
    _validate_model(model)

    body: dict[str, Any] = {
        "model": model,
        "prompt": prompt,
        "stream": stream,
    }
    body.update(kwargs)

    if stream:
        return _generate_stream(body, timeout=timeout)
    else:
        resp = _post_json("/api/generate", body, timeout=timeout)
        return resp.get("response", "")


def _generate_stream(body: dict, *, timeout: int) -> Generator[str, None, None]:
    """Internal generator for streaming generate responses."""
    for obj in _post_stream("/api/generate", body, timeout=timeout):
        token = obj.get("response", "")
        if token:
            yield token


def get_model_info(model: str) -> dict:
    """Get detailed information about a specific model.

    Queries POST /api/show to retrieve model metadata including
    parameters, template, system prompt, and license.

    Args:
        model: Ollama model name, e.g. "llama3:latest".

    Returns:
        Dict with model details. Typical keys:
            - modelfile (str): The Modelfile content
            - parameters (str): Model parameters
            - template (str): Prompt template
            - details (dict): Architecture info (family, parameter_size, etc.)

    Raises:
        ConnectionError: If Ollama is not running.
        ValueError: If the model is not found.
        RuntimeError: On other Ollama API errors.

    Example:
        >>> info = get_model_info("llama3:latest")
        >>> print(info.get("details", {}).get("parameter_size"))
    """
    _ensure_ollama()
    _validate_model(model)

    return _post_json("/api/show", {"name": model})
