"""
Local Bridge — Calls Ollama and other local LLM servers.

Makes HTTP calls to Ollama at localhost:11434. Supports both blocking and
streaming modes. Handles common failure cases: Ollama not running, model
not found, timeout.

All stdlib — no external packages.
"""

from __future__ import annotations

import http.client
import json
import logging
import re
import socket
import time
import urllib.request
from dataclasses import dataclass, field
from typing import Any, Generator, Optional

log = logging.getLogger("brynq.runtime.local_bridge")

# ---------------------------------------------------------------------------
# Data types
# ---------------------------------------------------------------------------

OLLAMA_HOST = "localhost"
OLLAMA_PORT = 11434
_MODEL_NAME_RE = re.compile(r"^[a-z0-9][a-z0-9:._-]*$")


def _validate_model_name(model: str) -> None:
    """Validate that a model name matches the expected pattern.

    Raises ValueError if the name is empty or contains invalid characters.
    """
    if not model or not _MODEL_NAME_RE.match(model):
        raise ValueError(
            f"Invalid model name '{model}'. "
            f"Model names must match the pattern ^[a-z0-9][a-z0-9:._-]*$"
        )
OLLAMA_BASE = f"http://{OLLAMA_HOST}:{OLLAMA_PORT}"
DEFAULT_TIMEOUT = 120  # seconds — large models can be slow on first load

# ---------------------------------------------------------------------------
# Ollama model list cache
# ---------------------------------------------------------------------------

_model_cache: dict[str, Any] = {}  # keys: "models" (list[str]), "timestamp" (float)
_MODEL_CACHE_TTL = 60  # seconds


def _get_cached_models(base_url: str = OLLAMA_BASE, timeout: int = 10) -> list[str]:
    """Return Ollama model names, using a 60-second cache.

    Fetches /api/tags and caches the result. Returns the cached list if
    it's less than 60 seconds old. Invalidates the cache on any error.
    """
    global _model_cache
    now = time.monotonic()
    if (
        _model_cache
        and _model_cache.get("base_url") == base_url
        and now - _model_cache.get("timestamp", 0) < _MODEL_CACHE_TTL
    ):
        return _model_cache["models"]

    try:
        url = f"{base_url}/api/tags"
        req = urllib.request.Request(url, method="GET", headers={"Accept": "application/json"})
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = json.loads(resp.read().decode("utf-8"))
        models = [m.get("name", "") for m in data.get("models", []) if m.get("name")]
        _model_cache = {"models": models, "timestamp": now, "base_url": base_url}
        return models
    except Exception:
        _model_cache = {}
        raise


def invalidate_model_cache() -> None:
    """Clear the Ollama model list cache."""
    global _model_cache
    _model_cache = {}


@dataclass
class BridgeResponse:
    """Standardised response from any bridge (local, cloud, or agent).

    Attributes:
        output: The full text response.
        success: Whether the call completed without error.
        error: Error message if the call failed, else None.
        elapsed_seconds: Wall-clock time for the call.
        token_estimate: Rough token count of the output.
        model: The model that was actually used.
        backend: The backend identifier (e.g. "ollama", "claude").
        metadata: Extra backend-specific metadata (e.g. usage stats).
    """

    output: str = ""
    success: bool = True
    error: Optional[str] = None
    elapsed_seconds: float = 0.0
    token_estimate: int = 0
    model: str = ""
    backend: str = ""
    metadata: dict = field(default_factory=dict)


def probe_ollama(base_url: str) -> bool:
    """Check if Ollama is reachable at the given base URL.

    Uses the cached /api/tags result when available.
    Returns True if the model list can be fetched, False otherwise.
    """
    try:
        _get_cached_models(base_url, timeout=5)
        return True
    except Exception:
        return False


def _estimate_tokens(text: str) -> int:
    """Rough token count: words * 1.3."""
    if not text or not text.strip():
        return 0
    import math

    return int(math.ceil(len(text.split()) * 1.3))


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


_CONNECTION_RETRIES = 3
_RETRY_DELAY = 2  # seconds

# ---------------------------------------------------------------------------
# Connection keepalive — reuse HTTP connection across blocking calls
# ---------------------------------------------------------------------------

_keepalive_conn: http.client.HTTPConnection | None = None


def _get_conn(timeout: int = DEFAULT_TIMEOUT) -> http.client.HTTPConnection:
    """Return a reusable HTTP connection to Ollama, creating one if needed."""
    global _keepalive_conn
    if _keepalive_conn is not None and _keepalive_conn.sock is not None:
        _keepalive_conn.timeout = timeout
        return _keepalive_conn
    # Close stale reference if any
    if _keepalive_conn is not None:
        try:
            _keepalive_conn.close()
        except Exception:
            pass
    _keepalive_conn = http.client.HTTPConnection(
        OLLAMA_HOST, OLLAMA_PORT, timeout=timeout
    )
    return _keepalive_conn


def _discard_conn() -> None:
    """Discard the cached connection after an error."""
    global _keepalive_conn
    if _keepalive_conn is not None:
        try:
            _keepalive_conn.close()
        except Exception:
            pass
        _keepalive_conn = None


def close_connection() -> None:
    """Close the persistent Ollama connection. Safe to call multiple times."""
    _discard_conn()


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

    Reuses a persistent HTTP connection for keepalive across calls.
    """
    data = json.dumps(body).encode("utf-8")
    headers = {"Content-Type": "application/json", "Accept": "application/json"}

    last_exc: Exception | None = None
    for attempt in range(_CONNECTION_RETRIES):
        conn = _get_conn(timeout)
        try:
            conn.request("POST", path, body=data, headers=headers)
            resp = conn.getresponse()
        except (ConnectionRefusedError, ConnectionResetError) as exc:
            _discard_conn()
            last_exc = exc
            if attempt < _CONNECTION_RETRIES - 1:
                log.warning(
                    "Ollama connection refused (attempt %d/%d), retrying in %ds...",
                    attempt + 1, _CONNECTION_RETRIES, _RETRY_DELAY,
                )
                time.sleep(_RETRY_DELAY)
                continue
            raise ConnectionError(
                f"Cannot reach Ollama at {OLLAMA_BASE}. "
                f"Is Ollama running? Start it with: ollama serve\n"
                f"Detail: {exc}"
            ) from exc
        except http.client.HTTPException as exc:
            # Connection in bad state (e.g. CannotSendRequest) — discard and retry
            _discard_conn()
            last_exc = exc
            if attempt < _CONNECTION_RETRIES - 1:
                continue
            raise ConnectionError(
                f"Cannot reach Ollama at {OLLAMA_BASE}. "
                f"Is Ollama running? Start it with: ollama serve\n"
                f"Detail: {exc}"
            ) from exc
        except (socket.timeout, OSError) as exc:
            _discard_conn()
            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")
            raise RuntimeError(
                f"Ollama returned HTTP {resp.status} for POST {path}: {detail}"
            )

        raw = resp.read().decode("utf-8")
        return json.loads(raw) if raw.strip() else {}

    # Should not reach here, but handle defensively
    raise ConnectionError(
        f"Cannot reach Ollama at {OLLAMA_BASE} after {_CONNECTION_RETRIES} attempts. "
        f"Is Ollama running? Start it with: ollama serve\n"
        f"Detail: {last_exc}"
    ) from last_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.
    """
    body["stream"] = True
    payload = json.dumps(body).encode("utf-8")

    last_exc: Exception | None = None
    conn: http.client.HTTPConnection | None = None
    resp: http.client.HTTPResponse | None = None
    for attempt in range(_CONNECTION_RETRIES):
        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()
            break
        except ConnectionRefusedError as exc:
            last_exc = exc
            if attempt < _CONNECTION_RETRIES - 1:
                log.warning(
                    "Ollama connection refused (attempt %d/%d), retrying in %ds...",
                    attempt + 1, _CONNECTION_RETRIES, _RETRY_DELAY,
                )
                time.sleep(_RETRY_DELAY)
                continue
            raise ConnectionError(
                f"Cannot reach Ollama at {OLLAMA_BASE}. "
                f"Is Ollama running? Start it with: ollama serve\n"
                f"Detail: {exc}"
            ) from exc
        except (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
    else:
        raise ConnectionError(
            f"Cannot reach Ollama at {OLLAMA_BASE} after {_CONNECTION_RETRIES} attempts. "
            f"Is Ollama running? Start it with: ollama serve\n"
            f"Detail: {last_exc}"
        ) from last_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:
                if buffer.strip():
                    try:
                        yield json.loads(buffer.decode("utf-8"))
                    except json.JSONDecodeError:
                        pass
                break
            buffer += chunk
            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()


# ---------------------------------------------------------------------------
# Local Bridge
# ---------------------------------------------------------------------------


def _get_available_models(base_url: str = OLLAMA_BASE, timeout: int = 10) -> list[str]:
    """GET /api/tags and return the list of model names (cached)."""
    return _get_cached_models(base_url, timeout=timeout)


def call_ollama(
    model: str,
    prompt: str,
    *,
    system_prompt: str = "",
    temperature: Optional[float] = None,
    top_p: Optional[float] = None,
    max_tokens: Optional[int] = None,
    timeout: int = DEFAULT_TIMEOUT,
    base_url: str = OLLAMA_BASE,
) -> BridgeResponse:
    """Call an Ollama model with a model existence check.

    Before sending the chat request, queries /api/tags to verify the model
    exists. Raises ValueError listing available models if it does not.
    """
    # --- Model name validation ---
    _validate_model_name(model)

    # --- Model existence check ---
    available = _get_available_models(base_url, timeout=10)
    # Match against full name and short name (without :tag)
    short_names = [n.split(":")[0] for n in available]
    if model not in available and model not in short_names:
        raise ValueError(
            f"Model '{model}' is not available in Ollama. "
            f"Available models: {', '.join(available) if available else '(none)'}. "
            f"Pull it with: ollama pull {model}"
        )

    # --- Build request ---
    messages: list[dict[str, str]] = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": prompt})

    body: dict[str, Any] = {"model": model, "messages": messages, "stream": False}
    options: dict[str, Any] = {}
    if temperature is not None:
        options["temperature"] = temperature
    if top_p is not None:
        options["top_p"] = top_p
    if max_tokens is not None:
        options["num_predict"] = max_tokens
    if options:
        body["options"] = options

    # --- Call Ollama ---
    response = BridgeResponse(model=model, backend="ollama")
    t0 = time.monotonic()
    try:
        resp = _post_json("/api/chat", body, timeout=timeout)
        output = resp.get("message", {}).get("content", "")
        if not isinstance(output, str) or not output.strip():
            log.warning(
                "Ollama model '%s' returned empty or non-string content",
                model,
            )
        response.output = output
        response.success = True
        response.token_estimate = _estimate_tokens(output)
        if "eval_count" in resp:
            response.metadata["eval_count"] = resp["eval_count"]
        if "prompt_eval_count" in resp:
            response.metadata["prompt_eval_count"] = resp["prompt_eval_count"]
    except ConnectionError as exc:
        response.success = False
        response.error = f"Ollama unavailable: {exc}"
        log.error("call_ollama connection error: %s", exc)
    except RuntimeError as exc:
        response.success = False
        response.error = f"Ollama error: {exc}"
        log.error("call_ollama runtime error: %s", exc)
    except Exception as exc:
        response.success = False
        response.error = f"Unexpected error: {exc}"
        log.error("call_ollama unexpected error: %s", exc, exc_info=True)
    finally:
        response.elapsed_seconds = time.monotonic() - t0
        if response.elapsed_seconds > 30:
            log.warning(
                "Ollama model '%s' took %.1fs to respond",
                model,
                response.elapsed_seconds,
            )
    return response


# ---------------------------------------------------------------------------
# LocalBridge class
# ---------------------------------------------------------------------------


class LocalBridge:
    """Bridge to local LLM servers (Ollama).

    Detects whether Ollama is running, validates model availability,
    and makes chat completion requests.
    """

    def __init__(
        self,
        host: str = OLLAMA_HOST,
        port: int = OLLAMA_PORT,
        timeout: int = DEFAULT_TIMEOUT,
    ) -> None:
        self.host = host
        self.port = port
        self.base_url = f"http://{host}:{port}"
        self.timeout = timeout

    def close(self) -> None:
        """Close the persistent HTTP connection to Ollama."""
        close_connection()

    def is_available(self) -> bool:
        """Check whether Ollama is running and reachable."""
        try:
            req = urllib.request.Request(f"{self.base_url}/", method="GET")
            with urllib.request.urlopen(req, timeout=5) as resp:
                return resp.status == 200
        except Exception:
            return False

    def list_models(self) -> list[str]:
        """Return names of all models available in Ollama (cached)."""
        try:
            return _get_cached_models(self.base_url, timeout=10)
        except Exception:
            return []

    def call(
        self,
        model: str,
        prompt: str,
        config: Optional[dict] = None,
    ) -> BridgeResponse:
        """Make a blocking chat completion call to Ollama.

        Args:
            model: Ollama model name (e.g. "llama3", "mistral").
            prompt: The user prompt to send.
            config: Optional config dict with keys like "temperature",
                    "max_tokens", "system_prompt".

        Returns:
            BridgeResponse with the model's output or error details.
        """
        config = config or {}
        _validate_model_name(model)
        response = BridgeResponse(model=model, backend="ollama")
        t0 = time.monotonic()

        try:
            # Check Ollama is running
            if not self.is_available():
                raise ConnectionError(
                    f"Ollama is not running at {self.base_url}. "
                    "Start it with: ollama serve"
                )

            # Build messages
            messages: list[dict[str, str]] = []
            system_prompt = config.get("system_prompt")
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            messages.append({"role": "user", "content": prompt})

            # Build request body
            body: dict[str, Any] = {
                "model": model,
                "messages": messages,
                "stream": False,
            }

            # Apply optional parameters
            options: dict[str, Any] = {}
            if "temperature" in config:
                options["temperature"] = config["temperature"]
            if "top_p" in config:
                options["top_p"] = config["top_p"]
            if "max_tokens" in config:
                options["num_predict"] = config["max_tokens"]
            if options:
                body["options"] = options

            # Make the call
            resp = _post_json("/api/chat", body, timeout=self.timeout)
            output = resp.get("message", {}).get("content", "")
            if not isinstance(output, str) or not output.strip():
                log.warning(
                    "Ollama model '%s' returned empty or non-string content",
                    model,
                )

            response.output = output
            response.success = True
            response.token_estimate = _estimate_tokens(output)

            # Extract usage metadata if available
            if "eval_count" in resp:
                response.metadata["eval_count"] = resp["eval_count"]
            if "prompt_eval_count" in resp:
                response.metadata["prompt_eval_count"] = resp["prompt_eval_count"]

        except ConnectionError as exc:
            response.success = False
            response.error = f"Ollama unavailable: {exc}"
            log.error("Local bridge connection error: %s", exc)

        except RuntimeError as exc:
            response.success = False
            error_str = str(exc)
            # Detect model-not-found errors
            if "not found" in error_str.lower() or "404" in error_str:
                response.error = (
                    f"Model '{model}' not found in Ollama. "
                    f"Pull it with: ollama pull {model}"
                )
            else:
                response.error = f"Ollama error: {exc}"
            log.error("Local bridge runtime error: %s", exc)

        except Exception as exc:
            response.success = False
            response.error = f"Unexpected error: {exc}"
            log.error("Local bridge unexpected error: %s", exc, exc_info=True)

        finally:
            response.elapsed_seconds = time.monotonic() - t0
            if response.elapsed_seconds > 30:
                log.warning(
                    "Ollama model '%s' took %.1fs to respond",
                    model,
                    response.elapsed_seconds,
                )

        return response

    def call_streaming(
        self,
        model: str,
        prompt: str,
        config: Optional[dict] = None,
    ) -> Generator[str, None, None]:
        """Make a streaming chat completion call to Ollama.

        Yields text chunks as they arrive. If an error occurs, yields
        a single chunk with the error message prefixed by "[ERROR] ".

        Args:
            model: Ollama model name.
            prompt: The user prompt.
            config: Optional config dict.

        Yields:
            Text chunks from the model's response.
        """
        config = config or {}
        _validate_model_name(model)

        if not self.is_available():
            yield f"[ERROR] Ollama is not running at {self.base_url}. Start it with: ollama serve"
            return

        messages: list[dict[str, str]] = []
        system_prompt = config.get("system_prompt")
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        body: dict[str, Any] = {
            "model": model,
            "messages": messages,
        }

        options: dict[str, Any] = {}
        if "temperature" in config:
            options["temperature"] = config["temperature"]
        if "top_p" in config:
            options["top_p"] = config["top_p"]
        if "max_tokens" in config:
            options["num_predict"] = config["max_tokens"]
        if options:
            body["options"] = options

        try:
            for obj in _post_stream("/api/chat", body, timeout=self.timeout):
                content = obj.get("message", {}).get("content", "")
                if content:
                    yield content
        except (ConnectionError, RuntimeError) as exc:
            yield f"[ERROR] {exc}"
        except Exception as exc:
            yield f"[ERROR] Unexpected: {exc}"
