prole/knoe/core/ollama_scan.py

"""Ollama network scan helpers.

This module contains the shared subprocess + parsing logic used by UI screens.
It intentionally has no GUI dependencies so it can be reused from multiple
screens without creating circular imports.
"""

from __future__ import annotations

import os
import re
import socket
import subprocess
import threading
from pathlib import Path
from typing import Callable

from knoe.core.env import DEFAULT_OLLAMA_PORT


def _is_ip(name: str) -> bool:
    if not name:
        return False
    return bool(re.match(r"^(?:\d{1,3}\.){3}\d{1,3}$", name)) or ":" in name


def _canon_host(name: str) -> str:
    n = (name or "").strip().lower()
    if n in {"localhost", "127.0.0.1", "::1", "k3d.local"}:
        return "127.0.0.1"
    try:
        return socket.gethostbyname(n)
    except Exception:
        return n


def _label_rank(name: str) -> int:
    n = (name or "").strip().lower()
    if n == "k3d.local":
        return 0
    if not _is_ip(n) and n not in {"localhost"}:
        return 1
    if n == "localhost":
        return 2
    return 3


def parse_ollama_scan_stdout(stdout_lines: list[str]) -> list[dict]:
    """Parse stdout lines from `init_ollama.sh scan`.

    Expected format per line: `host<TAB>port<TAB>model1,model2,...`.

    Returns rows: `{host: str, port: str, models: list[str]}`.
    """
    rows: list[dict] = []
    for line in stdout_lines:
        line = (line or "").strip()
        if not line or line.startswith("#"):
            continue
        parts = [p.strip() for p in line.split("\t")]
        if len(parts) < 2:
            continue
        host = parts[0]
        port = parts[1] or DEFAULT_OLLAMA_PORT
        model_list: list[str] = []
        if len(parts) > 2 and parts[2]:
            model_list = [
                m.strip()
                for m in parts[2].split(",")
                if m.strip() and m.strip() != "-"
            ]
        rows.append({"host": host, "port": str(port), "models": model_list})

    # Deduplicate by canonical host:port while preferring hostname labels.
    if not rows:
        return []

    merged: dict[tuple[str, str], dict] = {}
    for r in rows:
        disp_host = r.get("host", "")
        port = str(r.get("port", "") or DEFAULT_OLLAMA_PORT)
        canon = _canon_host(disp_host)
        key = (canon, port)
        cur = merged.get(key)
        if not cur:
            merged[key] = {
                "host": disp_host,
                "port": port,
                "models": list(r.get("models", []) or []),
            }
        else:
            existing = set(cur.get("models", []) or [])
            for m in r.get("models", []) or []:
                if m not in existing:
                    cur["models"].append(m)
                    existing.add(m)
            old_label = cur.get("host", "")
            if _label_rank(disp_host) < _label_rank(old_label):
                cur["host"] = disp_host

    return sorted(merged.values(), key=lambda d: (d.get("host") or ""))


def run_ollama_scan(
    project_root: Path,
    on_progress_line: Callable[[str], None] | None = None,
) -> tuple[int, list[dict], str]:
    """Run `etc/init_ollama.sh scan` and return `(rc, rows, notice)`.

    Progress is emitted on `stderr` and streamed to `on_progress_line`.
    Results are emitted on `stdout` and parsed into rows.
    """
    notice = ""
    rows: list[dict] = []
    script_path = project_root / "etc" / "init_ollama.sh"
    if not script_path.exists():
        notice = f"Missing script: {script_path}"
        if on_progress_line:
            try:
                on_progress_line(notice)
            except Exception:
                pass
        return 127, [], notice

    env = os.environ.copy()
    env["PROLE_HOME"] = str(project_root)
    env["PROLE_SERVICE"] = str(project_root)

    proc = subprocess.Popen(
        ["bash", str(script_path), "scan"],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=env,
        text=True,
        bufsize=1,
    )

    stdout_lines: list[str] = []

    def _read_stderr():
        if not proc.stderr:
            return
        for line in proc.stderr:
            line = line.rstrip("\n")
            if not line:
                continue
            if on_progress_line:
                try:
                    on_progress_line(line)
                except Exception:
                    pass

    stderr_thread = threading.Thread(target=_read_stderr, daemon=True)
    stderr_thread.start()

    if proc.stdout:
        for line in proc.stdout:
            line = line.rstrip("\n")
            if line:
                stdout_lines.append(line)

    proc.wait()
    stderr_thread.join(timeout=5)
    rc = int(proc.returncode or 0)

    if rc == 0:
        rows = parse_ollama_scan_stdout(stdout_lines)
        if rows:
            notice = f"Detected {len(rows)} server(s)."
        else:
            notice = "No Ollama servers detected."
    else:
        notice = "Scan failed. See console for details."

    if on_progress_line and notice:
        try:
            on_progress_line(notice)
        except Exception:
            pass

    return rc, rows, notice