#!/usr/bin/env bash
# u/complete: AI code completion. Reads prefix from stdin, prints completion to stdout.
set -euo pipefail

PREFIX_MAX=4000
SUFFIX_MAX=1000

MOUNT=${OLLIE:-$HOME/mnt/ollie}

cmd=${0##*/}
die() { printf '%s: %s\n' "$cmd" "$*" >&2; exit 1; }

for _arg; do
    [ "$_arg" = "-help" ] && cat <<'EOF' && exit 0
usage: complete [-model MODEL] [-backend BACKEND] [-file NAME] [-suffix TEXT]
       echo "prefix" | complete -file main.go -suffix "rest of code"

Reads prefix from stdin, prints raw completion to stdout.
Cost/info lines (:: prefix) go to stderr.

Flags:
  -model MODEL      model (default: $OLLIE_COMPLETE_MODEL)
  -backend BACKEND  backend (default: $OLLIE_COMPLETE_BACKEND)
  -file NAME        filename for language context
  -suffix TEXT      code after the cursor
  -help             show this help
EOF
done

model=${OLLIE_COMPLETE_MODEL:-}
backend=${OLLIE_COMPLETE_BACKEND:-}
file=
suffix=

while [ $# -gt 0 ]; do
    case $1 in
        -model)   model=$2;   shift ;;
        -backend) backend=$2; shift ;;
        -file)    file=$2;    shift ;;
        -suffix)  suffix=$2;  shift ;;
        -help)    ;;
        -*) die "unknown flag: $1" ;;
        *)  break ;;
    esac
    shift
done

[ -z "$model" ]   && die "model required: set OLLIE_COMPLETE_MODEL or pass -model"
[ -z "$backend" ] && die "backend required: set OLLIE_COMPLETE_BACKEND or pass -backend"

prefix=$(cat)

# Truncate to context budget
[ ${#prefix} -gt $PREFIX_MAX ] && prefix="${prefix: -$PREFIX_MAX}"
[ ${#suffix} -gt $SUFFIX_MAX ] && suffix="${suffix:0:$SUFFIX_MAX}"

# Find or create a copilot session for the current working directory.
find_or_create_session() {
    local target_cwd="$1"
    local id kv deadline

    id="$(printf '%s' "$target_cwd" | cksum | awk '{print $1}')-copilot"

    if [ -d "$MOUNT/s/$id" ]; then
        local state
        state=$(cat "$MOUNT/s/$id/state" 2>/dev/null) || true
        case "$state" in
            failed*|error*)
                printf 'kill\n' > "$MOUNT/s/$id/ctl" 2>/dev/null || true
                sleep 0.1
                ;;
            *)
                printf '%s\n' "$id"
                return 0
                ;;
        esac
    fi

    kv="name=${id}
cwd=${target_cwd}
agent=copilot
backend=${backend}
model=${model}"
    printf '%s\n' "$kv" > "$MOUNT/s/new" 2>/dev/null || true

    deadline=$(( $(date +%s) + 5 ))
    until [ -d "$MOUNT/s/$id" ] || [ "$(date +%s)" -ge "$deadline" ]; do sleep 0.05; done
    [ -d "$MOUNT/s/$id" ] || die "timeout waiting for copilot session"

    printf '%s\n' "$id"
}

sid=$(find_or_create_session "$PWD")
B="$MOUNT/s/$sid"

file_hint=
[ -n "$file" ] && file_hint=" in ${file}"

printf '%s\n' "Complete the code at the cursor${file_hint}. Continue from the prefix to finish the current function, block, or logical unit. Output ONLY the raw text to insert at the cursor position. Do NOT wrap the output in markdown code fences or backticks. No explanation, no preamble.

<prefix>
${prefix}
</prefix>
<suffix>
${suffix}
</suffix>" > "$B/prompt"

# Wait for idle
while true; do
    state=$(cat "$B/statewait" 2>/dev/null | tr -d '[:space:]') || break
    [ "$state" = "idle" ] && break
done

# Read result from offset, then clear context for next completion
offset=$(cat "$B/offset" 2>/dev/null || echo 0)
result=$(tail -c +$((offset + 1)) "$B/chat" 2>/dev/null)
# Strip markdown code fences if the model wrapped the output.
result=$(printf '%s' "$result" | sed '/^```[a-zA-Z]*$/d; /^```$/d')
# Print info lines (:: prefix) to stderr; keep only code.
printf '%s\n' "$result" | grep '^:: ' >&2 || true
result=$(printf '%s' "$result" | grep -v '^:: ')
echo "clear" > "$B/ctl"

printf '%s' "$result"