#!/usr/bin/env python3
"""
cluster.py — multi-node Ollama orchestrator.

Three potential nodes:
  - pc        : localhost (this PC, GPU, 10 models, fast but laptop-up only)
  - atlas-01  : <YOUR_VPS_IP>:11434 (Hetzner VPS, CPU, 5 models, always-on)
  - met       : <YOUR_PC_NODE_IP>:11434 (MET desktop, secondary capacity)

Routing strategy:
  - For "always_on" tasks (overnight, scheduled, server-side): prefer atlas-01
  - For "fast" tasks (interactive, jarvis, code completion): prefer pc
  - For "embed": prefer pc (nomic on both, pc is faster)
  - For "vision" (qwen2.5vl): pc only (atlas-01 doesn't have it)
  - Fallback chain: if preferred unreachable → try next available node that has the model

Public API:
    ask(prompt, task='general', node_pref=None, max_tokens=512) -> {ok, text, node, model, elapsed_s}
    embed(text) -> {ok, vector, node, model}
    nodes_status() -> [{name, reachable, latency_ms, models}]
    pick_node(task, model=None) -> (node_name, model_name)
"""
from __future__ import annotations
import json, time, urllib.request, urllib.error, socket

NODES = {
    'pc':       {'host': '127.0.0.1',     'port': 11434, 'always_on': False, 'fast': True,  'has_gpu': True},
    'atlas-01': {'host': '<YOUR_VPS_IP>',  'port': 11434, 'always_on': True,  'fast': False, 'has_gpu': False},
    # FIX 2026-05-26: was <YOUR_PC_NODE_IP> (which is PC's own tailnet IP).
    # MET-Desktop actually lives at <YOUR_DESKTOP_NODE_IP>.
    'met':      {'host': '<YOUR_DESKTOP_NODE_IP>','port': 11434, 'always_on': False, 'fast': True,  'has_gpu': False},
}

TASK_MODEL = {
    'chat':       'hermes3:8b',
    'voice':      'llama3.2:3b',
    'fast':       'llama3.2:1b',
    'code':       'qwen2.5-coder:7b',
    'code_fast':  'qwen2.5-coder:1.5b',
    'reason':     'phi3:medium',
    'general':    'qwen2.5:7b',
    'general_fast':'qwen2.5:1.5b',
    'embed':      'nomic-embed-text',
    'vision':     'qwen2.5vl:7b',
}

# Per-node model availability — populated lazily on first nodes_status() call.
_NODE_MODELS_CACHE = {'pc': None, 'atlas-01': None, 'met': None}
_CACHE_AT = {'pc': 0, 'atlas-01': 0, 'met': 0}
_CACHE_TTL = 300  # 5 min

def _http(url: str, body: dict | None = None, timeout=8):
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(url, data=data, method='POST' if body else 'GET')
    if body: req.add_header('Content-Type', 'application/json')
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read().decode())

def _check_node(name: str) -> dict:
    """Probe a node — returns {reachable, latency_ms, models, err?}."""
    if name not in NODES: return {'reachable': False, 'err': 'unknown node'}
    n = NODES[name]
    t0 = time.time()
    try:
        # Fast TCP check first
        with socket.create_connection((n['host'], n['port']), timeout=2): pass
        # Then fetch tags
        r = _http(f"http://{n['host']}:{n['port']}/api/tags", timeout=5)
        models = [m['name'] for m in r.get('models', [])]
        return {'reachable': True, 'latency_ms': int((time.time() - t0) * 1000), 'models': models}
    except (socket.error, urllib.error.URLError, urllib.error.HTTPError) as e:
        return {'reachable': False, 'latency_ms': int((time.time() - t0) * 1000), 'err': str(e)[:60]}

def nodes_status(force=False) -> dict:
    """Return per-node {reachable, models, latency_ms}. Cached 5 min."""
    out = {}
    now = time.time()
    for name in NODES:
        if not force and _NODE_MODELS_CACHE[name] is not None and now - _CACHE_AT[name] < _CACHE_TTL:
            out[name] = _NODE_MODELS_CACHE[name]
        else:
            s = _check_node(name)
            _NODE_MODELS_CACHE[name] = s
            _CACHE_AT[name] = now
            out[name] = s
    return out

def pick_node(task: str = 'general', model: str | None = None, prefer: str | None = None) -> tuple[str, str] | None:
    """Decide (node, model) for a task. Returns None if no node available."""
    if not model: model = TASK_MODEL.get(task, TASK_MODEL['general'])
    status = nodes_status()

    # Build preference order
    order = []
    if prefer and prefer in NODES: order.append(prefer)
    # task-driven defaults
    if task in ('embed',) or task.endswith('_fast') or task == 'fast' or task == 'voice':
        order += ['pc', 'atlas-01', 'met']
    elif task == 'vision':
        order += ['pc']  # only PC has vision model
    else:
        order += ['pc', 'atlas-01', 'met']
    # dedupe preserving order
    seen = set(); order = [n for n in order if not (n in seen or seen.add(n))]

    for name in order:
        s = status.get(name, {})
        if not s.get('reachable'): continue
        models = s.get('models', [])
        # exact match
        if model in models: return (name, model)
        # match-by-prefix (qwen2.5-coder:7b matches qwen2.5-coder:7b-instruct-q4_K_M etc.)
        for m in models:
            if m.startswith(model.split(':')[0]): return (name, m)
    return None

def ask(prompt: str, task: str = 'general', node_pref: str | None = None,
        max_tokens: int = 512, temperature: float = 0.4, timeout: int = 120) -> dict:
    pick = pick_node(task, prefer=node_pref)
    if not pick: return {'ok': False, 'err': 'no node has the required model'}
    node, model = pick
    n = NODES[node]
    t0 = time.time()
    try:
        r = _http(f"http://{n['host']}:{n['port']}/api/generate", {
            'model': model, 'prompt': prompt, 'stream': False,
            'options': {'temperature': temperature, 'num_predict': max_tokens},
        }, timeout=timeout)
        return {'ok': True, 'text': r.get('response', '').strip(),
                'node': node, 'model': model,
                'elapsed_s': round(time.time() - t0, 2)}
    except Exception as e:
        return {'ok': False, 'err': str(e)[:120], 'node': node, 'model': model,
                'elapsed_s': round(time.time() - t0, 2)}

def embed(text: str, node_pref: str | None = None) -> dict:
    pick = pick_node('embed', prefer=node_pref)
    if not pick: return {'ok': False, 'err': 'no embed model'}
    node, model = pick
    n = NODES[node]
    t0 = time.time()
    try:
        r = _http(f"http://{n['host']}:{n['port']}/api/embeddings", {
            'model': model, 'prompt': text,
        }, timeout=30)
        return {'ok': True, 'vector': r.get('embedding', []),
                'node': node, 'model': model,
                'elapsed_s': round(time.time() - t0, 2)}
    except Exception as e:
        return {'ok': False, 'err': str(e)[:120]}

if __name__ == '__main__':
    import sys
    cmd = sys.argv[1] if len(sys.argv) > 1 else 'status'
    if cmd == 'status':
        import io
        try: sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
        except: pass
        s = nodes_status(force=True)
        for name, info in s.items():
            mark = 'OK ' if info.get('reachable') else 'FAIL'
            mc = len(info.get('models', []))
            lat = info.get('latency_ms', '?')
            err = info.get('err', '')
            print(f'  {mark} {name:10s} {mc:>2d} models  {lat:>4}ms  {err}')
    elif cmd == 'ask':
        task = sys.argv[2] if len(sys.argv) > 2 else 'general'
        prompt = ' '.join(sys.argv[3:]) or 'say OK only'
        r = ask(prompt, task=task)
        if r['ok']:
            print(f'[{r["node"]} / {r["model"]} / {r["elapsed_s"]}s]')
            print(r['text'])
        else:
            print('ERR:', r.get('err'))
    elif cmd == 'pick':
        task = sys.argv[2] if len(sys.argv) > 2 else 'general'
        prefer = sys.argv[3] if len(sys.argv) > 3 else None
        p = pick_node(task, prefer=prefer)
        print(p or 'no node available')