#!/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}. The prefix is already written. Output ONLY the new text that comes after the prefix — do NOT repeat any part of the prefix. Continue to finish the current function, block, or logical unit. No markdown code fences, no 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 and leaked XML tags if the model wrapped the output.
result=$(printf '%s' "$result" | sed '/^```[a-zA-Z]*$/d; /^```$/d; /^<\/\?\(prefix\|suffix\)>$/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"

# Strip prefix echo: if the result starts with a suffix of the prefix, remove it.
# Check progressively shorter tails of the prefix (up to 200 chars).
tail_max=200
[ ${#prefix} -lt $tail_max ] && tail_max=${#prefix}
i=$tail_max
while [ "$i" -gt 0 ]; do
    tail="${prefix: -$i}"
    case "$result" in
        "$tail"*)
            result="${result#"$tail"}"
            break
            ;;
    esac
    i=$((i - 1))
done

printf '%s' "$result"