import asyncio
import json
import logging
import subprocess
import time
import httpx
import os
import signal
from pathlib import Path
from typing import Any

from .nexus_config import settings, DEFAULT_CHAT_MODEL, DEFAULT_EMBED_MODEL
from .memory.store import store

OLLAMA_PORT = 11434

_log = logging.getLogger("nexus.ollama")

# --- Global Singleton Instance ---
_ollama_manager = None

# Bundled binary ships alongside the project; fall back to system PATH
_BUNDLED_OLLAMA = settings.project_root / "ollama" / "bin" / (
    "ollama.exe" if os.name == "nt" else "ollama"
)

# POSIX: detach the child into its own session so we can signal the whole group.
# Windows has no setsid/killpg — run the child normally and terminate() it.
_DETACH_KW = {} if os.name == "nt" else {"start_new_session": True}


def _ollama_bin() -> str:
    """Return path to the Ollama executable, preferring the bundled copy."""
    if _BUNDLED_OLLAMA.exists():
        return str(_BUNDLED_OLLAMA)
    return "ollama"


def _best_vulkan_device() -> tuple[int, str]:
    """
    Parse `vulkaninfo --summary` and return (device_index, device_name) for the
    best Vulkan compute device. Prefers discrete GPUs over integrated ones, and
    AMD/NVIDIA vendor IDs over Intel — so a Radeon is chosen over an Intel iGPU
    even when the iGPU appears first in the device list.

    Software rasterizers (llvmpipe/lavapipe, PHYSICAL_DEVICE_TYPE_CPU) are
    dropped outright: Mesa always advertises one, it is CPU inference wearing a
    GPU costume, and it is *slower* than the plain CPU backend because every
    tensor takes a detour through Vulkan. Returns (-1, "") when no real GPU is
    present so the caller falls back instead of pinning the rasterizer.
    """
    try:
        r = subprocess.run(
            ["vulkaninfo", "--summary"], capture_output=True, text=True, timeout=5,
        )
        if r.returncode != 0:
            return -1, ""

        devices: list[dict] = []
        current: dict = {}
        for line in r.stdout.splitlines():
            line = line.strip()
            if line.startswith("GPU") and line.endswith(":"):
                if current:
                    devices.append(current)
                raw_idx = line[3:-1]
                current = {
                    "index": int(raw_idx) if raw_idx.isdigit() else len(devices),
                    "name": "GPU",
                    "type": "",
                    "vendor": "",
                }
            elif "=" in line:
                key, _, val = line.partition("=")
                key, val = key.strip(), val.strip()
                if key == "deviceName":
                    current["name"] = val
                elif key == "deviceType":
                    current["type"] = val.upper()
                elif key == "vendorID":
                    current["vendor"] = val.lower()

        if current:
            devices.append(current)

        devices = [d for d in devices if "CPU" not in d["type"]]
        if not devices:
            return -1, ""

        def _score(d: dict) -> tuple:
            # Discrete beats everything; integrated is last resort
            type_score = 2 if "DISCRETE" in d["type"] else (0 if "INTEGRATED" in d["type"] else 1)
            # AMD (0x1002) and NVIDIA (0x10de) preferred over Intel (0x8086)
            vendor_score = 1 if any(v in d["vendor"] for v in ("0x1002", "0x10de")) else 0
            return (type_score, vendor_score)

        best = max(devices, key=_score)
        return best["index"], best["name"]

    except Exception:
        return -1, ""


def _detect_gpu_backend() -> tuple[str, dict]:
    """
    Probe available GPU compute backends.
    Returns (label, env_overrides) where env_overrides is merged into the
    Ollama subprocess environment before launch.
    Priority: CUDA > Vulkan > ROCm > CPU.
    """
    # NVIDIA CUDA — preferred when both GPU and CUDA drivers are present
    try:
        r = subprocess.run(
            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
            capture_output=True, text=True, timeout=5,
        )
        if r.returncode == 0:
            name = r.stdout.strip().splitlines()[0]
            _log.info("GPU backend: CUDA (%s)", name)
            return f"cuda ({name})", {}  # Ollama auto-detects CUDA
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    # Vulkan — works on AMD, Intel, and NVIDIA without a full CUDA/ROCm stack
    vulkan_ok = False
    try:
        r = subprocess.run(
            ["vulkaninfo", "--summary"], capture_output=True, text=True, timeout=5,
        )
        vulkan_ok = r.returncode == 0
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    if not vulkan_ok:
        # Fall back to checking for ICD loader files without the vulkaninfo tool
        icd_dirs = [
            Path("/usr/share/vulkan/icd.d"),
            Path("/etc/vulkan/icd.d"),
            Path(os.path.expanduser("~/.local/share/vulkan/icd.d")),
        ]
        try:
            vulkan_ok = any(p.is_dir() and any(p.iterdir()) for p in icd_dirs)
        except PermissionError:
            pass

    if vulkan_ok:
        idx, name = _best_vulkan_device()
        if idx >= 0:
            # Always pin to the selected device — without this, Ollama may use the Intel
            # iGPU's shared system RAM as "VRAM" for models that don't fit on discrete VRAM.
            env_overrides: dict = {"OLLAMA_VULKAN": "1", "GGML_VK_VISIBLE_DEVICES": str(idx)}
            _log.info("GPU backend: Vulkan device %d (%s)", idx, name)
            return f"vulkan ({name})", env_overrides
        # Vulkan loads but every device is a software rasterizer — no real GPU here.

    # AMD ROCm — fallback when Vulkan ICD is absent but ROCm stack is installed
    try:
        r = subprocess.run(
            ["rocm-smi", "--showproductname"], capture_output=True, text=True, timeout=5,
        )
        if r.returncode == 0:
            _log.info("GPU backend: ROCm")
            return "rocm", {}  # Ollama auto-detects ROCm
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    _log.info("GPU backend: CPU (no GPU acceleration detected)")
    return "cpu", {}


# Single source of model auto-selection preference. Ordered so small, GPU-fitting
# models come first (qwen2.5:3b fits a 4GB card and is a strong all-rounder); the
# tail differs by task. Prefix-matched against installed model names.
_MODEL_PREFERENCE = {
    "chat": ("qwen2.5:3b", "qwen2.5", "gemma3:1b", "gemma3", "phi3", "phi-3", "gemma2", "gemma"),
    "code": ("qwen2.5-coder", "qwen3-coder", "deepseek-coder", "codellama", "codegemma", "qwen2.5:3b", "qwen2.5", "gemma3", "phi3", "phi-3"),
}


def _preferred_model(models: list, preference) -> str | None:
    """First installed model whose name starts with a preference prefix."""
    for prefix in preference:
        for m in models:
            if m.lower().startswith(prefix.lower()):
                return m
    return None


def _chat_options(temperature: float | None, num_gpu: int | None, num_ctx: int | None = None) -> dict:
    """Assemble the Ollama `options` block from the knobs we expose.

    Returns an empty dict when nothing is set so callers can omit `options`
    entirely (preserving Ollama's defaults / auto behaviour).
    """
    opts: dict = {}
    if temperature is not None:
        opts["temperature"] = temperature
    if num_gpu is not None:
        opts["num_gpu"] = num_gpu
    if num_ctx:  # 0 / None -> let Ollama use the model default
        opts["num_ctx"] = num_ctx
    return opts


_THINK_OPEN = "<think>"
_THINK_CLOSE = "</think>"


def _partial_tag_tail(text: str, tag: str) -> int:
    """Length of the longest suffix of `text` that could be the start of `tag`.

    A tag can arrive split across stream chunks ("<thi" then "nk>"), so that much
    of the tail has to be held back rather than emitted.
    """
    for n in range(min(len(tag) - 1, len(text)), 0, -1):
        if text.endswith(tag[:n]):
            return n
    return 0


class ThinkStripper:
    """Removes a reasoning model's <think> spans from a token stream.

    Ollama routes reasoning into `message.thinking` only when asked to think.
    We ask for `think: False` because the reasoning is pure latency here — but
    deepseek-r1 and friends emit `<think>` inline in `content` anyway, so the
    whole internal monologue reached the chat window, closing tags and all.

    Two shapes show up in practice. A well-formed span is dropped whole. A
    *stray* closing tag with no opening — which is what actually shipped — is at
    least removed, so the user does not see a literal `</think>` in the reply.
    Text already streamed before it cannot be recalled; `strip_think()` handles
    that case properly for callers that have the complete message.
    """

    def __init__(self) -> None:
        self._buf = ""
        self._inside = False

    def feed(self, chunk: str) -> str:
        self._buf += chunk
        out: list[str] = []
        while True:
            if self._inside:
                end = self._buf.find(_THINK_CLOSE)
                if end == -1:
                    keep = _partial_tag_tail(self._buf, _THINK_CLOSE)
                    self._buf = self._buf[len(self._buf) - keep:] if keep else ""
                    break
                self._buf = self._buf[end + len(_THINK_CLOSE):]
                self._inside = False
                continue

            start = self._buf.find(_THINK_OPEN)
            stray = self._buf.find(_THINK_CLOSE)
            # A stray close before any open: drop the tag, keep going.
            if stray != -1 and (start == -1 or stray < start):
                out.append(self._buf[:stray])
                self._buf = self._buf[stray + len(_THINK_CLOSE):]
                continue
            if start == -1:
                keep = max(
                    _partial_tag_tail(self._buf, _THINK_OPEN),
                    _partial_tag_tail(self._buf, _THINK_CLOSE),
                )
                if keep:
                    out.append(self._buf[:len(self._buf) - keep])
                    self._buf = self._buf[len(self._buf) - keep:]
                else:
                    out.append(self._buf)
                    self._buf = ""
                break
            out.append(self._buf[:start])
            self._buf = self._buf[start + len(_THINK_OPEN):]
            self._inside = True
        return "".join(out)

    def flush(self) -> str:
        """Whatever is left once the stream ends. An unterminated <think> is
        reasoning that never closed, so it is dropped rather than shown."""
        rest = "" if self._inside else self._buf
        self._buf = ""
        return rest


def strip_think(text: str) -> str:
    """Remove reasoning from a complete message.

    Unlike the streaming case this sees everything, so an unmatched closing tag
    can be handled the way it was meant: every token before it was reasoning,
    and the answer is what follows.
    """
    if not text or _THINK_CLOSE not in text and _THINK_OPEN not in text:
        return text
    import re
    cleaned = re.sub(r"<think>.*?</think>", "", text, flags=re.S)
    if _THINK_CLOSE in cleaned:               # stray close: the answer follows it
        cleaned = cleaned.rsplit(_THINK_CLOSE, 1)[1]
    cleaned = re.sub(r"<think>.*\Z", "", cleaned, flags=re.S)  # never closed
    return cleaned.strip()


async def _raise_for_ollama(r: httpx.Response) -> None:
    """raise_for_status(), but say what Ollama actually said.

    Ollama answers every failure with {"error": "..."} — a model that isn't
    pulled, a request too large for VRAM, a cloud model retired upstream — and
    httpx's default message discards the body, leaving the user with:

        Client error '410 Gone' for url 'http://127.0.0.1:11434/api/chat'

    when the body held "glm-4.6 was retired at 2026-06-16". Same exception type
    as before so existing handlers are unaffected; only the message improves.
    """
    if r.is_success:
        return
    # A streamed response has no body loaded yet; reading it is what makes the
    # error message available at all.
    try:
        await r.aread()
    except Exception:
        pass
    detail = ""
    try:
        body = r.json()
        if isinstance(body, dict):
            detail = str(body.get("error") or "").strip()
    except Exception:
        detail = (r.text or "").strip()
    if not detail:
        r.raise_for_status()  # nothing to add — keep httpx's wording
    raise httpx.HTTPStatusError(
        f"Ollama {r.status_code} from {r.request.url.path}: {detail[:400]}",
        request=r.request,
        response=r,
    )


class OllamaManager:
    def __init__(self, runtime_dir=None):
        self.process = None
        self.running = False
        self._available = None      # see is_available()

        self.runtime_dir = Path(runtime_dir) if runtime_dir else settings.runtime_dir
        (self.runtime_dir / "logs").mkdir(parents=True, exist_ok=True)

        self.log_file = self.runtime_dir / "logs" / "ollama.log"

        # Resolved at construction so host changes in settings take effect
        self._api_base = settings.ollama_host.rstrip("/")

        # Model selection cache
        self._model_cache: dict = {}  # intent -> (model, monotonic_ts)

        # Per-model offloadable layer count cache (never changes for a model)
        self._layer_cache: dict[str, int] = {}

        # How long Ollama keeps the model resident between requests. Applied to
        # every chat/generate body so the model isn't reloaded on each message.
        # Overridden from persisted settings at startup. Falsy → omit (Ollama's
        # 5-minute default).
        self.keep_alive: str | None = "30m"

    def _apply_keep_alive(self, body: dict) -> dict:
        """Add `keep_alive` (a top-level Ollama field) to a request body when set."""
        if self.keep_alive:
            body["keep_alive"] = self.keep_alive
        return body

    async def warm(self, model: str | None = None, num_gpu: int | None = None) -> None:
        """Preload a model so the first request doesn't pay a cold load.
        An empty-prompt /api/generate is Ollama's documented preload. Pass the
        same `num_gpu` the real requests use, or the preloaded copy is placed
        differently and gets reloaded on first use. Best-effort: never raises,
        so a missing model or down server can't break startup."""
        try:
            model = model or await self.select_best_model()
            if not model:
                return
            body = {"model": model, "prompt": "", "stream": False}
            if num_gpu is not None:
                body["options"] = {"num_gpu": num_gpu}
            async with httpx.AsyncClient(timeout=120.0) as client:
                await client.post(
                    f"{self._api_base}/api/generate", json=self._apply_keep_alive(body),
                )
            _log.info("warm: preloaded model=%s num_gpu=%s keep_alive=%s", model, num_gpu, self.keep_alive)
        except Exception as e:
            _log.warning("warm: preload failed: %s", e)

    def _serve_env(self, gpu_env: dict) -> dict:
        """Environment for a `serve` we spawn. One definition - this was
        duplicated verbatim in two methods, so the Windows carve-out below had
        to be fixed in both places or the two paths would disagree."""
        env = os.environ.copy()
        # The bind value, not the connect value: a user who set 0.0.0.0 to reach
        # Ollama from another machine must still get a server that listens there.
        env["OLLAMA_HOST"] = settings.ollama_bind
        # Every platform uses the project's own model store. Windows used to be
        # exempt, because the installer pulled with a bare `ollama pull` into
        # %USERPROFILE%\.ollama and a NexusOS-spawned serve pointed elsewhere
        # would not have seen those models. The installer sets OLLAMA_MODELS for
        # its pulls now (and migrates an existing store), so the exemption just
        # meant Windows kept models somewhere `ncp models` could not see.
        env["OLLAMA_MODELS"] = str(settings.models_dir)
        env.update(gpu_env)
        return env

    def is_available(self):
        if not settings.manage_ollama:
            # A remote provider has no local executable to discover. Availability
            # means it is configured; is_running() performs the live probe.
            return True
        # ponytail: cached for the life of the process. This spawns a subprocess,
        # and /status calls it on every poll - the frontend polls continuously,
        # so it was a process spawn per tick to answer a question whose answer
        # cannot change without someone installing Ollama. Restart to re-detect.
        if self._available is None:
            bin_path = _ollama_bin()
            try:
                subprocess.run([bin_path, "--version"], capture_output=True,
                               check=True, timeout=5)
                self._available = True
            except Exception:
                self._available = False
        return self._available

    def is_running(self):
        # 0.5s, not 2s: this is a loopback request to a server that is either
        # listening or is not. Ollama ships OFF (the user presses Start AI), so
        # the "not running" path is the common one and every /status paid the
        # full 2s for it - which pushed /status past the 2s client timeout in
        # bin/nexus_window.py and made a healthy backend look dead.
        try:
            r = httpx.get(f"{self._api_base}/api/tags", timeout=0.5)
            return r.status_code == 200
        except httpx.RequestError:
            return False

    def start(self):
        if not settings.manage_ollama:
            _log.info("Remote Ollama is externally managed at %s", self._api_base)
            return self.is_running()
        if not self.is_available():
            _log.warning("Ollama not found at %s; skipping startup", _ollama_bin())
            return False

        if self.is_running():
            _log.info("Ollama already running (%s)", self._api_base)
            self.running = True
            return True

        try:
            backend, gpu_env = _detect_gpu_backend()
            _log.info("Starting Ollama service via %s (backend: %s)...", _ollama_bin(), backend)
            env = self._serve_env(gpu_env)

            with open(self.log_file, "w") as log:
                self.process = subprocess.Popen(
                    [_ollama_bin(), "serve"],
                    stdout=log,
                    stderr=subprocess.STDOUT,
                    env=env,
                    **_DETACH_KW,
                )

            for attempt in range(30):
                if self.is_running():
                    _log.info("Ollama service started (%s)", self._api_base)
                    self.running = True
                    return True
                time.sleep(1)
                if attempt % 5 == 0:
                    _log.info("Waiting for Ollama... (%ds)", attempt)

            _log.error("Ollama failed to start: timeout")
            return False

        except Exception as e:
            _log.exception("Ollama failed to start: %s", e)
            return False

    async def start_async(self):
        """Async-safe version of start() for use inside async startup handlers."""
        if not settings.manage_ollama:
            _log.info("Remote Ollama is externally managed at %s", self._api_base)
            return self.is_running()
        if not self.is_available():
            _log.warning("Ollama not found at %s; skipping startup", _ollama_bin())
            return False

        if self.is_running():
            _log.info("Ollama already running (%s)", self._api_base)
            self.running = True
            return True

        try:
            backend, gpu_env = _detect_gpu_backend()
            _log.info("Starting Ollama service via %s (backend: %s)...", _ollama_bin(), backend)
            env = self._serve_env(gpu_env)

            with open(self.log_file, "w") as log:
                self.process = subprocess.Popen(
                    [_ollama_bin(), "serve"],
                    stdout=log,
                    stderr=subprocess.STDOUT,
                    env=env,
                    **_DETACH_KW,
                )

            for attempt in range(30):
                if self.is_running():
                    _log.info("Ollama service started (%s)", self._api_base)
                    self.running = True
                    return True
                await asyncio.sleep(1)
                if attempt % 5 == 0:
                    _log.info("Waiting for Ollama... (%ds)", attempt)

            _log.error("Ollama failed to start: timeout")
            return False

        except Exception as e:
            _log.exception("Ollama failed to start: %s", e)
            return False

    def stop(self):
        if not settings.manage_ollama:
            _log.info("Not stopping externally managed Ollama at %s", self._api_base)
            return False
        # Terminate a server we spawned ourselves.
        if self.process:
            try:
                _log.info("Stopping Ollama service (owned process)...")
                if os.name == "nt":
                    self.process.terminate()
                else:
                    os.killpg(os.getpgid(self.process.pid), signal.SIGTERM)
                self.process.wait(timeout=10)
                _log.info("Ollama service stopped")
            except Exception:
                try:
                    if os.name == "nt":
                        self.process.kill()
                    else:
                        os.killpg(os.getpgid(self.process.pid), signal.SIGKILL)
                except Exception:
                    pass
            finally:
                self.process = None

        # Ollama may still be up because something else started it (e.g. the
        # Windows desktop app autostarts one). The manual Stop button should
        # still stop the AI, so kill any remaining server by name. Best-effort:
        # an elevated instance can't be killed from a user-level process, so log
        # rather than raise.
        if self.is_running():
            _log.info("Stopping externally-started Ollama...")
            try:
                if os.name == "nt":
                    for image in ("ollama app.exe", "ollama.exe"):
                        subprocess.run(
                            ["taskkill", "/F", "/T", "/IM", image],
                            capture_output=True, timeout=10,
                        )
                else:
                    subprocess.run(
                        ["pkill", "-f", "ollama serve"], capture_output=True, timeout=10,
                    )
            except Exception as e:
                _log.warning("external Ollama stop failed: %s", e)

        self.running = False

    def get_status(self):
        if self.is_running():
            return "running"
        elif self.is_available():
            return "available"
        else:
            return "unavailable"

    async def generate(
        self,
        prompt: str,
        model: str = DEFAULT_CHAT_MODEL,
        stream: bool = False,
        system: str = "",
        **kwargs
    ):
        start = time.perf_counter()

        _log.debug("generate model=%s system=%r prompt=%.120s", model, system, prompt)

        try:
            if stream:
                return self._stream(prompt=prompt, model=model, system=system, start=start)
            else:
                async with httpx.AsyncClient(timeout=300.0) as client:
                    r = await client.post(
                        f"{self._api_base}/api/generate",
                        json=self._apply_keep_alive({
                            "model": model,
                            "prompt": prompt,
                            "system": system,
                            "stream": False,
                        }),
                    )

                elapsed = time.perf_counter() - start
                _log.info("generate completed model=%s status=%d elapsed=%.3fs", model, r.status_code, elapsed)
                await _raise_for_ollama(r)
                return r.json().get("response", "")

        except Exception as e:
            elapsed = time.perf_counter() - start
            _log.exception("generate error after %.3fs: %s", elapsed, e)
            return None

    async def _stream(self, prompt: str, model: str, system: str, start: float):
        """
        Async generator that streams token chunks from Ollama.
        Yields string chunks as they arrive.
        """
        try:
            async with httpx.AsyncClient(timeout=300.0) as client:
                async with client.stream(
                    "POST",
                    f"{self._api_base}/api/generate",
                    json=self._apply_keep_alive({
                        "model": model,
                        "prompt": prompt,
                        "system": system,
                        "stream": True,
                    }),
                ) as response:
                    await _raise_for_ollama(response)
                    async for line in response.aiter_lines():
                        if not line.strip():
                            continue
                        try:
                            data = json.loads(line)
                            token = data.get("response", "")
                            if token:
                                yield token
                            if data.get("done", False):
                                elapsed = time.perf_counter() - start
                                _log.info("generate stream completed model=%s elapsed=%.3fs", model, elapsed)
                                break
                        except Exception:
                            continue

        except Exception as e:
            elapsed = time.perf_counter() - start
            _log.exception("generate stream error after %.3fs: %s", elapsed, e)
            return


    async def chat(
        self,
        messages: list,
        model: str = DEFAULT_CHAT_MODEL,
        stream: bool = False,
        temperature: float | None = None,
        num_gpu: int | None = None,
        think: bool = False,
        tools: list | None = None,
        num_ctx: int | None = None,
        **kwargs,
    ):
        """Multi-turn chat via /api/chat (accepts a messages array with roles).

        When `tools` is given (non-stream only), the request advertises them and
        the FULL message dict is returned (so the caller sees `tool_calls`);
        otherwise the response content string is returned as before.

        `think` toggles Qwen3-style reasoning. Default off: the hidden <think>
        block is pure latency for chat/memory. Ollama ignores it for models that
        don't support thinking.
        """
        start = time.perf_counter()
        try:
            if stream:
                return self._chat_stream(
                    messages=messages, model=model, temperature=temperature,
                    num_gpu=num_gpu, think=think, start=start, num_ctx=num_ctx,
                )
            else:
                body: dict = {"model": model, "messages": messages, "stream": False}
                body["think"] = think
                if tools:
                    body["tools"] = tools
                opts = _chat_options(temperature, num_gpu, num_ctx)
                if opts:
                    body["options"] = opts
                self._apply_keep_alive(body)
                async with httpx.AsyncClient(timeout=300.0) as client:
                    r = await client.post(f"{self._api_base}/api/chat", json=body)
                elapsed = time.perf_counter() - start
                await _raise_for_ollama(r)
                message = r.json().get("message", {})
                # A reasoning model puts its monologue in `content` even with
                # think off, so strip it before anyone reads the answer.
                if isinstance(message, dict) and message.get("content"):
                    message["content"] = strip_think(message["content"])
                # Tool callers need the whole message (tool_calls); others want content.
                return message if tools else message.get("content", "")
        except Exception as e:
            elapsed = time.perf_counter() - start
            _log.exception("chat error after %.3fs: %s", elapsed, e)
            return None

    async def embed(self, text: str, model: str = DEFAULT_EMBED_MODEL) -> list[float] | None:
        """Return an embedding vector for `text` via /api/embeddings.

        Returns None on any failure so callers can fall back to lexical search —
        a missing embedding model should never break chat or recall.
        """
        text = (text or "").strip()
        if not text:
            return None
        try:
            gpu_offload = store.get_settings().get("memory_gpu_offload", 0)
            num_gpu = await self.resolve_num_gpu(gpu_offload, model)
            body: dict[str, Any] = {"model": model, "prompt": text}
            if num_gpu is not None:
                body["options"] = {"num_gpu": num_gpu}
            async with httpx.AsyncClient(timeout=30.0) as client:
                r = await client.post(
                    f"{self._api_base}/api/embeddings",
                    json=body,
                )
                await _raise_for_ollama(r)
                vec = r.json().get("embedding")
                return vec if vec else None
        except Exception as e:
            _log.debug("embed failed (model=%s): %s", model, e)
            return None

    async def list_models(self) -> list[str]:
        """Return names of all locally installed Ollama models."""
        try:
            async with httpx.AsyncClient(timeout=5.0) as client:
                r = await client.get(f"{self._api_base}/api/tags")
                await _raise_for_ollama(r)
                return [m["name"] for m in r.json().get("models", [])]
        except Exception:
            return []

    async def select_best_model(self, intent: str = "chat") -> str:
        """Pick the preferred installed model for `intent` ('chat' or 'code'),
        falling back to any installed model, then DEFAULT_CHAT_MODEL. Cached ~60s per
        intent so rapid requests don't rebuild the model list each time.
        """
        now = time.monotonic()
        cached = self._model_cache.get(intent)
        if cached and (now - cached[1]) < 60:
            return cached[0]

        models = await self.list_models()
        pref = _MODEL_PREFERENCE.get(intent, _MODEL_PREFERENCE["chat"])
        best = _preferred_model(models, pref) or (models[0] if models else DEFAULT_CHAT_MODEL)

        self._model_cache[intent] = (best, now)
        return best

    def invalidate_model_cache(self):
        """Force next select_best_model() to re-query (e.g. after pull/delete)."""
        self._model_cache = {}

    async def get_model_layers(self, model: str) -> int | None:
        """Total offloadable layer count for `model` (repeating blocks + output layer).

        Used to turn a CPU/GPU offload percentage into an Ollama `num_gpu`
        value. Reads `<arch>.block_count` from /api/show and adds 1 for the
        non-repeating output layer (Ollama reports e.g. 33 layers for a model
        with block_count=32). Cached per-model since it never changes.
        Returns None if the count can't be determined, so callers fall back
        to Auto (no num_gpu override).
        """
        if model in self._layer_cache:
            return self._layer_cache[model]
        try:
            async with httpx.AsyncClient(timeout=10.0) as client:
                r = await client.post(f"{self._api_base}/api/show", json={"model": model})
            await _raise_for_ollama(r)
            info = r.json().get("model_info", {}) or {}
            block_count = next(
                (v for k, v in info.items() if k.endswith(".block_count")), None
            )
            layers = int(block_count) + 1 if block_count is not None else None
        except Exception as e:
            _log.warning("get_model_layers(%s) failed: %s", model, e)
            layers = None
        if layers:
            self._layer_cache[model] = layers
        return layers

    async def resolve_num_gpu(self, gpu_offload, model: str) -> int | None:
        """Convert a stored gpu_offload setting into an Ollama `num_gpu` value.

        `gpu_offload` is -1 for Auto (returns None → no override, Ollama auto-fits)
        or 0–100 for the percent of the model's layers to force onto the GPU
        (0 = all CPU/RAM). Layer count is model-specific, resolved from the live
        model. Returns None on anything unexpected so callers fall back to Auto.
        """
        try:
            pct = int(gpu_offload)
        except (TypeError, ValueError):
            return None
        if pct < 0:
            return None
        pct = min(pct, 100)
        layers = await self.get_model_layers(model)
        if not layers:
            return None
        return max(0, round(pct / 100 * layers))

    async def _chat_stream(self, messages: list, model: str, start: float,
                           temperature: float | None = None, num_gpu: int | None = None,
                           think: bool = False, num_ctx: int | None = None):
        """Async generator streaming tokens, then a final __meta__ stats sentinel."""
        try:
            body: dict = {"model": model, "messages": messages, "stream": True}
            body["think"] = think  # see chat(): reasoning off by default for speed
            opts = _chat_options(temperature, num_gpu, num_ctx)
            if opts:
                body["options"] = opts
            self._apply_keep_alive(body)
            async with httpx.AsyncClient(timeout=300.0) as client:
                async with client.stream(
                    "POST",
                    f"{self._api_base}/api/chat",
                    json=body,
                ) as response:
                    await _raise_for_ollama(response)
                    thinking = ThinkStripper()
                    async for line in response.aiter_lines():
                        if not line.strip():
                            continue
                        try:
                            data = json.loads(line)
                            token = data.get("message", {}).get("content", "")
                            if token:
                                token = thinking.feed(token)
                                if token:
                                    yield token
                            if data.get("done", False):
                                tail = thinking.flush()  # held-back partial tag
                                if tail:
                                    yield tail
                                elapsed = time.perf_counter() - start
                                _log.info("chat stream completed model=%s elapsed=%.3fs", model, elapsed)
                                eval_count = data.get("eval_count", 0)
                                eval_ns = data.get("eval_duration", 0)
                                tokens_per_s = round(eval_count / (eval_ns / 1e9), 1) if eval_ns else 0
                                stats = json.dumps({
                                    "model": model,
                                    "tokens": eval_count,
                                    "elapsed_s": round(elapsed, 2),
                                    "tokens_per_s": tokens_per_s,
                                })
                                yield f"__meta__{stats}"
                                break
                        except Exception:
                            continue
        except Exception as e:
            elapsed = time.perf_counter() - start
            _log.exception("chat stream error after %.3fs: %s", elapsed, e)
            raise


def initialize_ollama() -> OllamaManager:
    global _ollama_manager

    if _ollama_manager is None:
        manager = OllamaManager()

        if not manager.is_running():
            manager.start()

        if not manager.is_running():
            raise RuntimeError("Ollama API is not reachable after start().")

        _ollama_manager = manager

    return _ollama_manager


async def initialize_ollama_async() -> OllamaManager:
    """Async-safe initializer — uses asyncio.sleep so the event loop stays live."""
    global _ollama_manager

    if _ollama_manager is None:
        manager = OllamaManager()

        if not manager.is_running():
            await manager.start_async()

        if not manager.is_running():
            raise RuntimeError("Ollama API is not reachable after start().")

        _ollama_manager = manager

    return _ollama_manager


def get_ollama_manager() -> OllamaManager:
    global _ollama_manager
    if _ollama_manager is None:
        _ollama_manager = OllamaManager()
    return _ollama_manager


def shutdown_ollama() -> None:
    global _ollama_manager
    if _ollama_manager is not None:
        _ollama_manager.stop()
        _ollama_manager = None
