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

PREFIX_MAX=12000
SUFFIX_MAX=1000

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

for _arg; do
    [ "$_arg" = "-help" ] && cat <<'HELP' && 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.

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
  -context TEXT     additional context (LSP symbols, completions, hover)
  -cwd DIR          working directory
  -help             show this help
HELP
done

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

while [ $# -gt 0 ]; do
    case $1 in
        -model)   model=$2;   shift ;;
        -backend) backend=$2; shift ;;
        -file)    file=$2;    shift ;;
        -suffix)  suffix=$2;  shift ;;
        -context) context=$2; shift ;;
        -cwd)     cwd=$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}"

# JSON-escape a string.
jsonesc() {
    printf '%s' "$1" | sed 's/\\/\\\\/g; s/"/\\"/g; s/\t/\\t/g' | tr '\n' '\1' | sed 's/\x01/\\n/g'
}

# Build JSON request for /complete endpoint
json='{"prefix":"'$(jsonesc "$prefix")'"'
json="$json,\"suffix\":\"$(jsonesc "$suffix")\""
[ -n "$file" ]    && json="$json,\"file\":\"$(jsonesc "$file")\""
[ -n "$context" ] && json="$json,\"context\":\"$(jsonesc "$context")\""
json="$json,\"cwd\":\"$(jsonesc "${cwd:-$PWD}")\""
json="$json}"

# Write to /complete and read response (net/dns pattern)
result=$(printf '%s' "$json" | 9p rdwr ollie/complete) || die "complete failed (is olliesrv running?)"

# Strip markdown code fences if the model wrapped the output.
result=$(printf '%s' "$result" | sed '/^```[a-zA-Z]*$/d; /^```$/d')

# Strip prefix echo: if the result starts with a suffix of the prefix, remove it.
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"