#!/usr/bin/env bash
set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
ENV_FILE="${ROOT}/.env"
[[ -f "$ENV_FILE" ]] && source "$ENV_FILE"

: "${OLLAMA_URL:=http://127.0.0.1:11434}"
: "${EXPERT_MODEL:=jr-sql-expert}"

ts() { date -Is; }

log_dir="${ROOT}/logs"
mkdir -p "$log_dir"
log_file="${log_dir}/sqlai-$(date -I).log"

# Log everything (stdout+stderr)
exec > >(tee -a "$log_file") 2>&1

usage() {
  cat <<'EOF'
Usage:
  sqlai <mode> [--file path] [--text "free text"] [--no-metrics]

Modes:
  ask              General SQL Server 2022 Q&A (free text)
  analyze-tsql     Analyze a T-SQL query
  analyze-proc     Analyze a stored procedure
  analyze-view     Analyze a view
  analyze-plan     Analyze an execution plan (Showplan XML or text)
  utf8-migration   Create a UTF-8 migration plan (SQL Server 2022)

Input options:
  --text "..."     Provide input text directly (no pipe needed)
  --file path      Read input from file
  (no args)        Reads from STDIN

Other options:
  --no-metrics     Do not print metrics line

Examples:
  sqlai ask --text "How do I troubleshoot parameter sniffing in SQL Server 2022?"
  echo "SELECT 1;" | sqlai analyze-tsql
  sqlai analyze-plan --file showplan.xml
EOF
}

mode="${1:-}"
[[ -z "$mode" ]] && { usage; exit 1; }
shift || true

file=""
text=""
no_metrics="0"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --file) file="${2:-}"; shift 2;;
    --text) text="${2:-}"; shift 2;;
    --no-metrics) no_metrics="1"; shift 1;;
    -h|--help) usage; exit 0;;
    *) echo "[$(ts)] ERROR: Unknown arg: $1"; usage; exit 2;;
  esac
done

# Read input (priority: --text > --file > stdin)
input=""
input_src=""
if [[ -n "$text" ]]; then
  input="$text"
  input_src="text"
elif [[ -n "$file" ]]; then
  [[ -f "$file" ]] || { echo "[$(ts)] ERROR: file not found: $file"; exit 3; }
  input="$(cat "$file")"
  input_src="file:${file}"
else
  input="$(cat)"
  input_src="stdin"
fi

if [[ -z "${input//[[:space:]]/}" ]]; then
  echo "[$(ts)] ERROR: empty input (source=$input_src)"
  exit 4
fi

# Short instruction per mode (core policy is in Modelfile SYSTEM prompt)
case "$mode" in
  ask)
    instruction=$'Beantworte die Frage als SQL Server 2022 Experte.\n\nWichtig:\n- Wenn Kontext fehlt: keine Rückfragen stellen; stattdessen Annahmen offenlegen und Optionen (A/B/C) mit Vor-/Nachteilen geben.\n- Ergebnis immer strukturiert mit: Kurzfazit, Optionen, Risiken/Checks, Nächste Schritte.\n- Wenn sinnvoll: konkrete T-SQL/DDL Snippets und Checklisten liefern.'
    ;;
  analyze-tsql)
    instruction=$'Analysiere das folgende T-SQL (SQL Server 2022): Performance, SARGability, Datentypen, Joins/Predicates, Indizes/Stats, Parameter Sniffing. Gib konkrete Rewrite- und Index-Ideen.'
    ;;
  analyze-proc)
    instruction=$'Analysiere die folgende Stored Procedure (SQL Server 2022): Performance/Correctness, Transaktionen/Locking, Sniffing, Temp tables vs table variables, RBAR. Gib konkrete Refactorings.'
    ;;
  analyze-view)
    instruction=$'Analysiere die folgende View (SQL Server 2022): SARGability, Predicate Pushdown, Expand/Inlining, Aggregationen/Distinct/Union, UDF-Risiken. Gib konkrete Verbesserungen.'
    ;;
  analyze-plan)
    instruction=$'Analysiere den folgenden SQL Server Execution Plan (Showplan XML oder Text): Hotspots, Spills/Warnings, Memory Grants, Kardinalität, Fixes (Rewrite/Stats/Indexing vorsichtig/Sniffing Mitigation).'
    ;;
  utf8-migration)
    instruction=$'Erstelle einen Migrationsplan (SQL Server 2022) zur Umstellung auf UTF-8 (UTF-8 enabled collations _UTF8): Abhängigkeiten (PK/FK/Indexes/Computed/Triggers), Cutover, Rollback, Tests, konkrete DDL/Checklisten.'
    ;;
  *)
    echo "[$(ts)] ERROR: unknown mode: $mode"
    usage
    exit 5
    ;;
esac

req_id="$(date +%Y%m%d-%H%M%S)-$$-$RANDOM"

echo "================================================================================"
echo "[$(ts)] sqlai: REQUEST_START id=$req_id"
echo "[$(ts)] sqlai: MODE=$mode MODEL=$EXPERT_MODEL OLLAMA_URL=$OLLAMA_URL"
echo "[$(ts)] sqlai: INPUT_SOURCE=$input_src INPUT_BYTES=$(printf "%s" "$input" | wc -c)"

# Delimited markup
prompt=$(cat <<EOF
${instruction}

---BEGIN INPUT (${mode})---
${input}
---END INPUT (${mode})---
EOF
)

# Build JSON safely (no heredoc+herestring combos)
payload="$(
printf '%s' "$prompt" | python -c '
import json, os, sys
model=os.environ.get("EXPERT_MODEL","jr-sql-expert")
prompt=sys.stdin.read()
print(json.dumps({"model": model, "prompt": prompt, "stream": False}, ensure_ascii=False))
'
)"

resp_file="$(mktemp)"
http_code="$(
curl -sS -o "$resp_file" -w "%{http_code}" \
  -X POST "${OLLAMA_URL}/api/generate" \
  -H 'Content-Type: application/json' \
  --data-binary "$payload" \
  || true
)"

resp="$(cat "$resp_file")"
rm -f "$resp_file"

echo "[$(ts)] sqlai: HTTP_CODE=$http_code RESP_BYTES=$(printf "%s" "$resp" | wc -c) id=$req_id"

# Validate JSON
if ! printf '%s' "$resp" | python -c 'import json,sys; json.load(sys.stdin)' >/dev/null 2>&1; then
  echo "[$(ts)] sqlai: ERROR: response is not valid JSON id=$req_id"
  echo "[$(ts)] sqlai: RAW_RESPONSE_BEGIN id=$req_id"
  printf '%s\n' "$resp"
  echo "[$(ts)] sqlai: RAW_RESPONSE_END id=$req_id"
  echo "[$(ts)] sqlai: REQUEST_END id=$req_id status=error"
  exit 11
fi

# Extract error/response/metrics in one pass
extracted="$(
printf '%s' "$resp" | python -c '
import json,sys
obj=json.load(sys.stdin)
out={
  "error": obj.get("error"),
  "response": obj.get("response",""),
  "metrics": {
    "total_duration": obj.get("total_duration"),
    "load_duration": obj.get("load_duration"),
    "prompt_eval_count": obj.get("prompt_eval_count"),
    "prompt_eval_duration": obj.get("prompt_eval_duration"),
    "eval_count": obj.get("eval_count"),
    "eval_duration": obj.get("eval_duration"),
  }
}
print(json.dumps(out, ensure_ascii=False))
'
)"

error_msg="$(printf '%s' "$extracted" | python -c 'import json,sys; print((json.load(sys.stdin).get("error") or "").strip())')"
response_txt="$(printf '%s' "$extracted" | python -c 'import json,sys; print(json.load(sys.stdin).get("response") or "")')"

# HTTP != 200 is error
if [[ "$http_code" != "200" ]]; then
  echo "[$(ts)] sqlai: ERROR: non-200 HTTP_CODE=$http_code id=$req_id"
  if [[ -n "$error_msg" ]]; then
    echo "[$(ts)] sqlai: OLLAMA_ERROR=$error_msg id=$req_id"
  else
    echo "[$(ts)] sqlai: BODY_SNIPPET_BEGIN id=$req_id"
    printf '%s\n' "$resp" | head -n 120
    echo "[$(ts)] sqlai: BODY_SNIPPET_END id=$req_id"
  fi
  echo "[$(ts)] sqlai: REQUEST_END id=$req_id status=error"
  exit 12
fi

# API-level error
if [[ -n "$error_msg" ]]; then
  echo "[$(ts)] sqlai: ERROR: OLLAMA_ERROR=$error_msg id=$req_id"
  echo "[$(ts)] sqlai: REQUEST_END id=$req_id status=error"
  exit 13
fi

# Print answer
printf "\n%s\n\n" "$(printf "%s" "$response_txt" | sed 's/[[:space:]]*$//')"

# If response empty, dump JSON snippet to log
if [[ -z "${response_txt//[[:space:]]/}" ]]; then
  echo "[$(ts)] sqlai: WARN: empty response id=$req_id"
  echo "[$(ts)] sqlai: RAW_JSON_SNIPPET_BEGIN id=$req_id"
  printf '%s\n' "$resp" | head -n 200
  echo "[$(ts)] sqlai: RAW_JSON_SNIPPET_END id=$req_id"
fi

# Print metrics optionally
if [[ "$no_metrics" != "1" ]]; then
  metrics_line="$(printf '%s' "$extracted" | python -c 'import json,sys; print("METRICS="+json.dumps(json.load(sys.stdin)["metrics"], ensure_ascii=False))')"
  echo "$metrics_line"
fi

echo "[$(ts)] sqlai: REQUEST_END id=$req_id status=ok"
echo "================================================================================"