#!/usr/bin/env zsh
# llm-ask.sh — Thin CLI wrapper for oMLX. Pipe-friendly.
#
# Usage:
#   llm-ask.sh "What is the capital of France?"
#   echo "Explain this code" | llm-ask.sh
#   cat file.py | llm-ask.sh "What does this do?"
#   llm-ask.sh --max-tokens 512 "Summarise this" < doc.txt
#
# Options:
#   --max-tokens N     Max output tokens (default: 2048)
#   --no-stream        Wait for full response instead of streaming
#   --port N           oMLX port (default: 8090 via proxy, or 8080 direct)
#   --model NAME       Model name (default: from /v1/models)

set -euo pipefail

PORT=8090
MAX_TOKENS=2048
STREAM=true
MODEL=""
PROMPT=""

# ── Parse args ────────────────────────────────────────────────────────────────
while [[ $# -gt 0 ]]; do
  case "$1" in
    --max-tokens) MAX_TOKENS="$2"; shift 2 ;;
    --no-stream)  STREAM=false; shift ;;
    --port)       PORT="$2"; shift 2 ;;
    --model)      MODEL="$2"; shift 2 ;;
    --help|-h)
      sed -n '2,15p' "$0" | sed 's/^# \?//'
      exit 0 ;;
    *) PROMPT="$PROMPT $1"; shift ;;
  esac
done

PROMPT="${PROMPT# }"  # trim leading space

# ── Read stdin if available ───────────────────────────────────────────────────
STDIN_CONTENT=""
if [[ ! -t 0 ]]; then
  STDIN_CONTENT=$(cat)
fi

if [[ -z "$PROMPT" && -z "$STDIN_CONTENT" ]]; then
  echo "Usage: llm-ask.sh [options] \"prompt\"" >&2
  echo "       echo \"prompt\" | llm-ask.sh" >&2
  exit 1
fi

# Combine stdin + prompt
if [[ -n "$STDIN_CONTENT" && -n "$PROMPT" ]]; then
  FULL_PROMPT="${STDIN_CONTENT}

${PROMPT}"
elif [[ -n "$STDIN_CONTENT" ]]; then
  FULL_PROMPT="$STDIN_CONTENT"
else
  FULL_PROMPT="$PROMPT"
fi

# ── Resolve model ─────────────────────────────────────────────────────────────
if [[ -z "$MODEL" ]]; then
  MODEL=$(curl -sf "http://localhost:${PORT}/v1/models" \
    | python3 -c "import sys,json; print(json.load(sys.stdin)['data'][0]['id'])" 2>/dev/null \
    || echo "gemma-4-26b-a4b-it-4bit")
fi

# ── Escape prompt for JSON ────────────────────────────────────────────────────
ESCAPED=$(python3 -c "import json,sys; print(json.dumps(sys.stdin.read()))" <<< "$FULL_PROMPT")

# ── Call oMLX ─────────────────────────────────────────────────────────────────
PAYLOAD="{\"model\":\"${MODEL}\",\"messages\":[{\"role\":\"user\",\"content\":${ESCAPED}}],\"max_tokens\":${MAX_TOKENS},\"stream\":${STREAM}}"

if [[ "$STREAM" == "true" ]]; then
  curl -sf "http://localhost:${PORT}/v1/chat/completions" \
    -H "Content-Type: application/json" \
    -d "$PAYLOAD" \
    --no-buffer \
  | python3 -u -c "
import sys, json
for line in sys.stdin:
    line = line.strip()
    if not line.startswith('data:'):
        continue
    data = line[5:].strip()
    if data == '[DONE]':
        break
    try:
        chunk = json.loads(data)
        delta = chunk.get('choices', [{}])[0].get('delta', {})
        text = delta.get('content') or delta.get('reasoning_content') or ''
        if text:
            print(text, end='', flush=True)
    except Exception:
        pass
print()
"
else
  curl -sf "http://localhost:${PORT}/v1/chat/completions" \
    -H "Content-Type: application/json" \
    -d "$PAYLOAD" \
  | python3 -c "import sys,json; print(json.load(sys.stdin)['choices'][0]['message']['content'])"
fi