#!/bin/bash
# Modelfile authoring and model pull for the local LLM.
#
# Sourced by setup_local_llm.sh; split out to keep it under the 250-line cap.
# Sourced rather than run, so it inherits the caller's strict mode and
# the helper functions and variables defined above the source line.

# ── 4. Create Modelfile with developer system prompt ─────────────────────────
create_modelfile() {
    section "Creating custom model with developer system prompt"

    # Condense Claude preferences into a system prompt (~1 K tokens).
    # Note: hooks, skills, and Claude Code MCP tools do NOT port to local models.
    # This only captures code-style, persona, and workflow preferences.
    local modelfile
    modelfile="$(mktemp --suffix=.modelfile)"

    cat > "$modelfile" << MODELFILE_EOF
FROM ${MODEL}

PARAMETER num_ctx ${CTX_LEN}
PARAMETER temperature 0.6
PARAMETER top_p 0.95
PARAMETER top_k 20

SYSTEM """
You are a coding assistant for Krzysztof (kuhy) Rudnicki, a software developer
who maintains personal automation and app monorepos.

## About me
- Stack: Python, Bash, C/C++, TypeScript, Dart/Flutter, Go
- OS: Arch Linux, i3 window manager
- Key repos: testsAndMisc (automation), linux-configuration, praca_magisterska

## Code style
Python:
- Use from __future__ import annotations for forward refs
- Google docstring convention (triple-quoted, brief summary + Args/Returns)
- Absolute imports only, no relative imports
- Type hints required on every function signature
- Prefer typed dataclasses over plain dicts

Shell (bash):
- Always set -euo pipefail at the top
- Double-quote all variable expansions
- Avoid fork-heavy patterns: prefer /proc, /sys, bash builtins over \$(...)
  in tight loops or status-bar scripts (fork-storm anti-pattern)
- Use jq/yq for JSON/YAML, not grep/awk
- NEVER embed multi-line Python logic inside shell heredocs — put it in a
  separate .py file and invoke with python3 path/to/file.py

TypeScript:
- Target ES2022, pure ESM (no CommonJS require)
- No any — use unknown + narrowing
- Discriminated unions for state machines
- async/await + structured error handling

## Behaviour I expect
- Keep responses brief and direct — no filler preamble
- Ask before making high-impact changes (schema changes, deleting data)
- If you spot a bad approach, say so before complying
- Development order: implement → run to verify → confirm → tests
  (tests are always LAST, not first)
- Include only one short comment per non-obvious block; no inline summaries
  of what the code obviously does

## Testing rules
- 100% branch coverage required for Python packages
- Use pytest + coverage; never mock the filesystem in integration tests
- Shell scripts: shellcheck must pass with no suppressions

## Git workflow
- Commit directly to main (no branches for personal repo)
- Conventional commit format: type: brief description
- Always run pre-commit hooks before committing
"""
MODELFILE_EOF

    local custom_model="${MODEL/:/-}-dev"
    ollama create "$custom_model" -f "$modelfile"
    rm -f "$modelfile"
    ok "Custom model created: $custom_model (includes developer system prompt)"
    ok "Load '$custom_model' in Open WebUI for kuhy-personalised responses"
}

# ── 5. Pull base model ────────────────────────────────────────────────────────
pull_model() {
    section "Pulling ${MODEL}"

    # Check if already present
    if curl -sf "http://localhost:${OLLAMA_PORT}/api/tags" \
            | grep -q "\"${MODEL}\""; then
        ok "${MODEL} already downloaded"
        return
    fi

    warn "Downloading ${MODEL} — approx 20 GB. This will take a while."
    warn "Progress is shown live; do not interrupt."
    ollama pull "${MODEL}"
    ok "${MODEL} downloaded"
}

# ── 6. Set up Open WebUI via Docker ───────────────────────────────────────────
setup_webui() {
    section "Setting up Open WebUI"

    # Remove existing container so we can apply any config changes cleanly
    if docker ps -a --format '{{.Names}}' | grep -qx "$WEBUI_CONTAINER"; then
        warn "Removing existing $WEBUI_CONTAINER container..."
        docker rm -f "$WEBUI_CONTAINER"
    fi

    docker pull ghcr.io/open-webui/open-webui:main

    docker run -d \
        --name "$WEBUI_CONTAINER" \
        --restart always \
        -p "${WEBUI_PORT}:8080" \
        -v "${WEBUI_VOLUME}:/app/backend/data" \
        -e OLLAMA_BASE_URL="http://host.docker.internal:${OLLAMA_PORT}" \
        --add-host=host.docker.internal:host-gateway \
        ghcr.io/open-webui/open-webui:main

    ok "Open WebUI container started on port ${WEBUI_PORT}"
}

# ── 7. Enable Docker on boot ──────────────────────────────────────────────────
enable_autostart() {
    section "Enabling auto-start on boot"

    # Docker daemon must start for the container (restart=always) to come up
    sudo systemctl enable docker
    ok "docker.service enabled on boot"
    ok "Ollama is already enabled; Open WebUI container uses restart=always"
}

# ── 8. Verify: GPU utilisation + generation speed ─────────────────────────────
verify_gpu() {
    section "Verifying GPU utilisation and generation speed"

    # Warm up with a short generation and measure tok/s
    warn "Running a warm-up generation (checking GPU offload + speed)..."
    local start_ts end_ts elapsed toks_per_sec
    start_ts="$(date +%s%N)"

    local response
    response="$(curl -sf "http://localhost:${OLLAMA_PORT}/api/generate" \
        -d "{\"model\":\"${MODEL}\",\"prompt\":\"Reply with exactly: ready\",\"stream\":false}" \
        --max-time 120)"

    end_ts="$(date +%s%N)"
    elapsed=$(( (end_ts - start_ts) / 1000000 ))  # milliseconds

    if [[ -z "$response" ]]; then
        err "No response from Ollama — generation failed"
        return 1
    fi

    # Parse eval_count (output tokens) and eval_duration (nanoseconds) from response
    local eval_count eval_duration_ns
    eval_count="$(printf '%s' "$response" | grep -o '"eval_count":[0-9]*' | grep -o '[0-9]*' || echo 0)"
    eval_duration_ns="$(printf '%s' "$response" | grep -o '"eval_duration":[0-9]*' | grep -o '[0-9]*' || echo 1)"

    if (( eval_count > 0 && eval_duration_ns > 0 )); then
        toks_per_sec=$(( eval_count * 1000000000 / eval_duration_ns ))
        ok "Generation speed: ${toks_per_sec} tok/s (${eval_count} tokens in ${elapsed}ms)"
        if (( toks_per_sec < 10 )); then
            warn "Speed is below 10 tok/s — model may be partially CPU-offloaded"
            warn "Run: ollama ps  to check GPU layer count"
            warn "If GPU% < 100, try reducing CTX_LEN in this script or use a smaller quant"
        fi
    else
        warn "Could not parse tok/s from response; check manually with: ollama ps"
    fi

    # Check GPU offload via ollama ps
    warn "Checking GPU layer utilisation..."
    if command -v ollama &>/dev/null; then
        ollama ps 2>/dev/null | head -10 || true
    fi

    ok "Verify done — see ollama ps output above for GPU% column"
}

# ── 9. WebUI readiness check ──────────────────────────────────────────────────
verify_webui() {
    section "Waiting for Open WebUI to become ready"
    wait_for_url "http://localhost:${WEBUI_PORT}" "Open WebUI" 60
}

# ── Summary ───────────────────────────────────────────────────────────────────
print_summary() {
    local custom_model="${MODEL/:/-}-dev"
    printf '\n\033[1;32m'
    printf '═══════════════════════════════════════════════════════\n'
    printf ' Local LLM setup complete!\n'
    printf '═══════════════════════════════════════════════════════\n'
    printf ' Open WebUI:     http://localhost:%s\n' "$WEBUI_PORT"
    printf ' Ollama API:     http://localhost:%s\n' "$OLLAMA_PORT"
    printf ' Primary model:  %s  (%d K context)\n' "$MODEL" $(( CTX_LEN / 1024 ))
    printf ' Dev model:      %s  (with code-style system prompt)\n' "$custom_model"
    printf '\n'
    printf ' GPU fit notes:\n'
    printf '   Current context: %d K (safe for 24 GB VRAM)\n' $(( CTX_LEN / 1024 ))
    printf '   To try 24 K:    edit CTX_LEN=24576 in this script\n'
    printf '   To try 32 K:    edit CTX_LEN=32768 (risks partial CPU offload)\n'
    printf '   After changes:  re-run this script; check tok/s in verify step\n'
    printf '\n'
    printf ' To use with Claude Code:\n'
    printf '   Claude Code uses the Anthropic API — it cannot point at Ollama\n'
    printf '   directly (different wire format). Use Open WebUI for local-model\n'
    printf '   coding sessions, or ask about setting up a LiteLLM proxy.\n'
    printf '═══════════════════════════════════════════════════════\n'
    printf '\033[0m\n'
}

# ── Main ──────────────────────────────────────────────────────────────────────