#!/usr/bin/env bash
# ai-edit: AI-abistaja koodi ja tekstide toimetamiseks
# Usage:
# In Termux: pkg install curl jq
# In Debian/Ubuntu: sudo apt install curl jq
# Initial Indexing (Large text/codebase):
# ./ai-edit -f draft_essay.md --build-notes
#
# ./ai-edit -f draft_essay.md -t 40 "Are the transitions between arguments clear?"
# ./ai-edit -f handlers.go -r 45:90 -c schema.sql "Ensure SQL parameters match the schema types"
# ./ai-edit -f story.txt -t 50 -s ~/.prompts/fiction_editor.txt "Critique dialogue cadence"
# Interaktiivne režiim:
# ./ai-edit-termux.sh -f raamat.txt -i -x "korrasta peatükid"
#   -x / --exec lülitab sisse käskude režiimi

set -euo pipefail

GOOGLE_API_KEY="${GOOGLE_API_KEY:-}"
MODEL="gemma-4-26b-a4b-it"
API_URL="https://generativelanguage.googleapis.com/v1beta/models/${MODEL}:generateContent"

[[ -z "${GOOGLE_API_KEY:-}" ]] && { echo "Viga: GOOGLE_API_KEY puudub. export GOOGLE_API_KEY=..." >&2; exit 1; }
command -v jq >/dev/null 2>&1 || { echo "Viga: 'jq' on vajalik." >&2; exit 1; }

EXEC_START="<<<EXEC>>>"
EXEC_END="<<<END>>>"
EXEC_MODE=0

usage() {
  cat << EOF
Kasutus: $(basename "$0") -f <raamat> [valikud] [instruktsioon]

  -f, --file FAIL      Raamatu fail
  -i, --interactive    Interaktiivne vestlus (jälgib faili jooksvalt)
  -t, --tail [N]       Viimased N rida
  -r, --range M:N      Reavahemik
  -c, --context FAIL   Lisakontekst
  -s, --sys FAIL       Süsteemiviip
  -n, --notes FAIL     Märkmete fail (vaikimisi <fail>.notes)
  -x, --exec           Luba AI-l pakkuda shell käske ${EXEC_START}/${EXEC_END} vahel
  --build-notes        Sunni täis-indekseerimine
  -h, --help

Termux näide:
  pkg install curl jq termux-api
  export GOOGLE_API_KEY="..."
  ./ai-edit-termux.sh -f raamat.txt -i -x -s ~/.prompts/kirjanik.txt "korrasta dialoog"

Kui AI vastab:
  ${EXEC_START}
  sed -i 's/  \\+/ /g' raamat.txt
  ${EXEC_END}
... siis küsitakse kinnitust ja saab [y] käivitada või [c] kopeerida teise akna jaoks.
EOF
}

TARGET_FILE=""; NOTES_FILE=""; TAIL_LINES=""; LINE_RANGE=""; SYS_FILE=""
INTERACTIVE=0; BUILD_NOTES=0; CONTEXT_FILES=()
DEFAULT_SYS="You are a precise technical editor and creative writing assistant. Keep notes compact."

SYSTEM_PROMPT="$DEFAULT_SYS"

while [[ $# -gt 0 ]]; do
  case "$1" in
    -f|--file) TARGET_FILE="${2:-}"; shift 2 ;;
    -i|--interactive) INTERACTIVE=1; shift 1 ;;
    -t|--tail)
      if [[ $# -gt 1 && "${2:-}" =~ ^[0-9]+$ ]]; then TAIL_LINES="$2"; shift 2
      else TAIL_LINES=50; shift 1; fi ;;
    -r|--range) LINE_RANGE="${2:-}"; shift 2 ;;
    -c|--context) CONTEXT_FILES+=("${2:-}"); shift 2 ;;
    -s|--sys) SYS_FILE="${2:-}"; shift 2 ;;
    -n|--notes) NOTES_FILE="${2:-}"; shift 2 ;;
    -x|--exec) EXEC_MODE=1; shift 1 ;;
    --build-notes) BUILD_NOTES=1; shift 1 ;;
    -h|--help) usage; exit 0 ;;
    *) break ;;
  esac
done
USER_INSTRUCTION="$*"

[[ -z "$TARGET_FILE" ]] && { echo "Viga: -f kohustuslik" >&2; usage; exit 1; }
if [[ ! -f "$TARGET_FILE" ]]; then
  echo -n "Faili '$TARGET_FILE' ei leitud. Luua? [y/N]: " >&2
  read -r ans; [[ "$ans" =~ ^[Yy]$ ]] && touch "$TARGET_FILE" || exit 1
fi
[[ -z "$NOTES_FILE" ]] && NOTES_FILE="${TARGET_FILE}.notes"
STATE_FILE="${NOTES_FILE}.state"
EXEC_LOG="${TARGET_FILE}.exec.log"

if [[ -n "$SYS_FILE" && -f "$SYS_FILE" ]]; then
  SYSTEM_PROMPT=$(cat "$SYS_FILE")
fi

# Lisa exec juhis system prompti kui -x on sees ja kasutaja sys fail seda juba ei sisalda
if [[ "$EXEC_MODE" -eq 1 ]]; then
  if ! grep -q "EXEC" <<< "$SYSTEM_PROMPT"; then
    SYSTEM_PROMPT+=$'\n\n'"[KÄSKUDE REŽIIM - TERMUX]
Kui kasutaja palub teksti korraldada, vormindada, asendada, sorteerida - ära kirjuta ainult juttu, paku ka käivitatav shell käsk.
Vorming peab olema TÄPSELT:
${EXEC_START}
# üks või mitu shell käsku, kasuta faili nime: ${TARGET_FILE}
sed -i 's/midagi/midagi/g' ${TARGET_FILE}
${EXEC_END}
Reeglid:
- Kasuta ainult ohutuid tekstikäsklusi: sed, awk, grep, tr, sort, uniq, perl -pe, head, tail, wc
- Ära kasuta rm -rf, curl | bash, sudo, mkfs, dd vms.
- Seleta lühidalt enne käsku mida teed.
- Kui käsk on ohtlik, hoiata.
"
  fi
fi

get_mtime() {
  if stat -c %Y "$1" >/dev/null 2>&1; then stat -c %Y "$1"
  elif stat -f %m "$1" >/dev/null 2>&1; then stat -f %m "$1"
  else echo 0; fi
}

call_gemini() {
  local prompt="$1" sys_prompt="$2" payload response text
  payload=$(jq -n --arg sp "$sys_prompt" --arg p "$prompt" \
    '{system_instruction:{parts:[{text:$sp}]}, contents:[{role:"user",parts:[{text:$p}]}], generationConfig:{temperature:0.3}}')
  response=$(curl -s -X POST "${API_URL}?key=${GOOGLE_API_KEY}" -H "Content-Type: application/json" -d "$payload")
  text=$(echo "$response" | jq -r '[.candidates[0].content.parts[]? | select(.thought != true) | .text // empty] | join("")')
  [[ -z "$text" ]] && { echo "API Viga:" >&2; echo "$response" | jq -r '.error.message // .' >&2; return 1; }
  echo "$text"
}

# ---- EXEC käsu käsitlus ----
extract_exec_commands() {
  local reply="$1"
  local cmds=""
  # 1. Otsi <<<EXEC>>> ... <<<END>>>
  if grep -qF "$EXEC_START" <<< "$reply"; then
    cmds=$(sed -n "/$(printf '%s' "$EXEC_START" | sed 's/\[/\\[/g; s/\]/\\]/g; s/>/\\>/g; s/</\\</g')/,/$(printf '%s' "$EXEC_END" | sed 's/\[/\\[/g; s/\]/\\]/g; s/>/\\>/g; s/</\\</g')/p" <<< "$reply" | sed "1d;\$d")
  # 2. fallback: ```bash ... ``` või ```sh ... ```
  elif grep -q '```bash\|```sh' <<< "$reply"; then
    cmds=$(echo "$reply" | sed -n '/```\(bash\|sh\)/,/```/p' | sed '1d;$d')
  fi
  echo "$cmds"
}

handle_exec() {
  local reply="$1"
  local cmds
  cmds=$(extract_exec_commands "$reply")
  [[ -z "$(echo "$cmds" | tr -d '[:space:]')" ]] && return 0

  echo -e "\n--- AI pakub käsku ---" >&2
  echo "$cmds" >&2
  echo "----------------------" >&2
  echo "$cmds" >> "$EXEC_LOG"
  echo "# $(date) - $TARGET_FILE" >> "$EXEC_LOG"

  # Küsi kinnitust
  while true; do
    echo -n "[y] käivita, [c] kopeeri lõikelauale, [e] redigeeri, [n] vahele? " >&2
    read -r choice
    case "$choice" in
      y|Y)
        echo "[*] Käivitan..." >&2
        # Käivita eraldi bashis, et set -e ei tapaks põhiprogrammi
        bash -c "$cmds"
        echo "[+] Käsk lõpetas koodiga $?" >&2
        # Pärast käivitust kontrolli kas raamat muutus
        echo "$(wc -l < "$TARGET_FILE")" > "$STATE_FILE"
        touch "$NOTES_FILE"
        break
        ;;
      c|C)
        if command -v termux-clipboard-set >/dev/null 2>&1; then
          echo "$cmds" | termux-clipboard-set
          echo "[+] Kopeeritud lõikelauale (termux-clipboard-set). Kleebi teises aknas." >&2
        else
          echo "[*] termux-api puudub, kopeeri käsitsi ülevalt." >&2
          echo "$cmds"
        fi
        break
        ;;
      e|E)
        tmpf=$(mktemp)
        echo "$cmds" > "$tmpf"
        ${EDITOR:-nano} "$tmpf"
        edited=$(cat "$tmpf")
        echo "Redigeeritud käsk:"
        echo "$edited" >&2
        echo -n "Käivita redigeeritud? [y/N]: " >&2
        read -r yn
        if [[ "$yn" =~ ^[Yy]$ ]]; then bash -c "$edited"; fi
        rm -f "$tmpf"
        break
        ;;
      n|N|"")
        echo "[*] Jäetud vahele." >&2
        break
        ;;
    esac
  done
}

# ---- Märkmete loogika ----
build_notes_full() {
  local total_lines=$(wc -l < "$TARGET_FILE") start=1 end chunk existing prompt updated
  echo "[*] Täis-indekseerimine $total_lines rida -> $NOTES_FILE" >&2
  rm -f "$NOTES_FILE"; touch "$NOTES_FILE"
  while [[ "$start" -le "$total_lines" ]]; do
    end=$((start+149)); [[ "$end" -gt "$total_lines" ]] && end="$total_lines"
    echo "   $start..$end" >&2
    chunk=$(sed -n "${start},${end}p" "$TARGET_FILE")
    existing=""; [[ -s "$NOTES_FILE" ]] && existing=$(cat "$NOTES_FILE")
    prompt="Current Notes: ${existing:-[None]}\nNew Chunk ${start}-${end}:\n\`\`\`\n${chunk}\n\`\`\`\nTask: Integrate into notes. Return ONLY notes."
    updated=$(call_gemini "$prompt" "You generate structural outlines.") || { start=$((end+1)); continue; }
    [[ -n "$updated" ]] && echo "$updated" > "$NOTES_FILE"
    start=$((end+1))
  done
  echo "$total_lines" > "$STATE_FILE"
}

build_notes_incremental() {
  local last=0 cur
  [[ -f "$STATE_FILE" ]] && last=$(cat "$STATE_FILE" 2>/dev/null || echo 0)
  [[ ! "$last" =~ ^[0-9]+$ ]] && last=0
  cur=$(wc -l < "$TARGET_FILE")
  [[ "$cur" -le "$last" ]] && { echo "$cur" > "$STATE_FILE"; touch "$NOTES_FILE"; return 0; }
  if [[ $((cur-last)) -gt 400 ]]; then build_notes_full; return; fi
  echo "[*] Auto-uuendus +$((cur-last)) rida ($last->$cur)" >&2
  local new_chunk=$(sed -n "$((last+1)),${cur}p" "$TARGET_FILE")
  local existing=$(cat "$NOTES_FILE" 2>/dev/null || echo "[None]")
  local prompt="Notes: ${existing}\nNew lines ${last}-${cur}:\n\`\`\`\n${new_chunk}\n\`\`\`\nIntegrate new into notes. Return ONLY updated notes."
  local updated=$(call_gemini "$prompt" "You generate structural outlines.") || return 1
  [[ -n "$updated" ]] && echo "$updated" > "$NOTES_FILE"
  echo "$cur" > "$STATE_FILE"; touch "$NOTES_FILE"
}

ensure_notes_up_to_date() {
  if [[ ! -f "$NOTES_FILE" ]]; then
    echo "[*] Märkmeid pole, loon..." >&2
    echo -e "# Teose märkmed\n- Teema:\n- Tegelased:\n- Süžee:" > "$NOTES_FILE"
    build_notes_full; return
  fi
  [[ ! -s "$NOTES_FILE" ]] && { build_notes_full; return; }
  local mb=$(get_mtime "$TARGET_FILE") mn=$(get_mtime "$NOTES_FILE")
  if [[ "$mb" -gt "$mn" ]]; then build_notes_incremental
  else
    if [[ -f "$STATE_FILE" ]]; then
      local last=$(cat "$STATE_FILE"); local cur=$(wc -l < "$TARGET_FILE")
      [[ "$cur" -gt "$last" ]] && build_notes_incremental
    fi
  fi
}

[[ "$BUILD_NOTES" -eq 1 ]] && { build_notes_full; exit 0; }
ensure_notes_up_to_date

USER_PAYLOAD=""
[[ -f "$NOTES_FILE" ]] && USER_PAYLOAD+="## Notes ${NOTES_FILE}
$(cat "$NOTES_FILE")

"
for ctx in "${CONTEXT_FILES[@]:-}"; do [[ -f "$ctx" ]] && USER_PAYLOAD+="## Context $ctx
\`\`\`
$(cat "$ctx")
\`\`\`

"; done
USER_PAYLOAD+="## File ${TARGET_FILE}
"
if [[ -n "$TAIL_LINES" ]]; then
  USER_PAYLOAD+="Last $TAIL_LINES lines
\`\`\`
$(tail -n "$TAIL_LINES" "$TARGET_FILE")
\`\`\`

"
elif [[ -n "$LINE_RANGE" ]]; then
  USER_PAYLOAD+="(Lines $LINE_RANGE)
\`\`\`
$(sed -n "${LINE_RANGE/:/,}p" "$TARGET_FILE")
\`\`\`

"
else
  tot=$(wc -l < "$TARGET_FILE")
  if [[ "$tot" -gt 350 ]]; then USER_PAYLOAD+="(Full $tot lines, showing last 120)
\`\`\`
$(tail -n 120 "$TARGET_FILE")
\`\`\`

"
  else USER_PAYLOAD+="(Full)
\`\`\`
$(cat "$TARGET_FILE")
\`\`\`

"
  fi
fi
[[ -z "$USER_INSTRUCTION" ]] && USER_INSTRUCTION="Vaata üle tekst, kontrolli sidusust märkmetega."
USER_PAYLOAD+="## Task
${USER_INSTRUCTION}"

if [[ "$INTERACTIVE" -eq 1 ]]; then
  HISTORY=$(jq -n --arg p "$USER_PAYLOAD" '[{"role":"user","parts":[{"text":$p}]}]')
  LAST_MTIME=$(get_mtime "$TARGET_FILE"); LAST_LINES=$(wc -l < "$TARGET_FILE")

  call_gemini_chat() {
    local payload=$(jq -n --arg sp "$SYSTEM_PROMPT" --argjson contents "$HISTORY" \
      '{system_instruction:{parts:[{text:$sp}]}, contents:$contents, generationConfig:{temperature:0.3}}')
    local resp=$(curl -s -X POST "${API_URL}?key=${GOOGLE_API_KEY}" -H "Content-Type: application/json" -d "$payload")
    local reply=$(echo "$resp" | jq -r '[.candidates[0].content.parts[]? | select(.thought != true) | .text // empty] | join("")')
    if [[ -n "$reply" ]]; then
      echo -e "\n$reply"
      HISTORY=$(echo "$HISTORY" | jq --arg r "$reply" '. + [{"role":"model","parts":[{"text":$r}]}]')
      [[ "$EXEC_MODE" -eq 1 ]] && handle_exec "$reply"
    else
      echo -e "\nAPI viga:" >&2; echo "$resp" | jq . >&2
      HISTORY=$(echo "$HISTORY" | jq 'del(.[-1])')
    fi
  }

  check_changed() {
    local cm=$(get_mtime "$TARGET_FILE") cl=$(wc -l < "$TARGET_FILE")
    if [[ "$cm" -ne "$LAST_MTIME" || "$cl" -ne "$LAST_LINES" ]]; then
      echo -e "\n[*] Fail muutus $LAST_LINES->$cl rida, uuendan märkmeid..." >&2
      build_notes_incremental || true
      LAST_MTIME=$cm; LAST_LINES=$cl
      local nt=$(tail -n 80 "$TARGET_FILE")
      local sys="[SÜSTEEM: Raamat salvestati. Uued märkmed: $(cat "$NOTES_FILE" | head -n 30) ... Viimane tekst: $nt ]"
      HISTORY=$(echo "$HISTORY" | jq --arg q "$sys" '. + [{"role":"user","parts":[{"text":$q}]}]')
      HISTORY=$(echo "$HISTORY" | jq '. + [{"role":"model","parts":[{"text":"[OK uuendasin]"}]}]')
    fi
  }

  call_gemini_chat
  echo -e "\n--- Vestlus avatud (jälgimine sees, -x sees). exit väljub ---"
  while true; do
    check_changed
    echo -n -e "\nSina > "; read -r fu || break
    [[ -z "$fu" ]] && continue
    check_changed
    case "$fu" in
      exit|q|quit) break ;;
      \>\>n*)
        note_extra=$(echo "$fu" | sed -E 's/^\>\>n[[:space:]]*//')
        last=$(echo "$HISTORY" | jq -r 'last(.[] | select(.role=="model")) .parts[0].text // empty')
        pr="Tee bullet täiendus märkmetesse. Fookus: $note_extra Kontekst: $last"
        if nr=$(call_gemini "$pr" "concise note-taker"); then echo -e "\n$nr" >> "$NOTES_FILE"; echo "[+] $NOTES_FILE" >&2; fi
        continue ;;
      \>\>*)
        df=$(echo "$fu" | sed 's/^>>[[:space:]]*//'); [[ -z "$df" ]] && df="$TARGET_FILE"
        last=$(echo "$HISTORY" | jq -r 'last(.[] | select(.role=="model")) .parts[0].text // empty')
        [[ -n "$last" ]] && echo -e "\n$last" >> "$df" && echo "[+] $df" >&2
        continue ;;
    esac
    HISTORY=$(echo "$HISTORY" | jq --arg q "$fu" '. + [{"role":"user","parts":[{"text":$q}]}]')
    call_gemini_chat
  done
else
  out=$(call_gemini "$USER_PAYLOAD" "$SYSTEM_PROMPT")
  echo "$out"
  [[ "$EXEC_MODE" -eq 1 ]] && handle_exec "$out"
fi