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

usage() {
  cat <<'EOF'
Usage: refactor_conn_aliasing.sh [OPTIONS] FILE...
  --dry-run    Show diff without modifying files
  --help       Show usage
EOF
}

DRY_RUN=false
FILES=()

while [[ $# -gt 0 ]]; do
  case "$1" in
    --dry-run) DRY_RUN=true; shift ;;
    --help)    usage; exit 0 ;;
    -*)        echo "Unknown option: $1" >&2; usage >&2; exit 1 ;;
    *)         FILES+=("$1"); shift ;;
  esac
done

if [[ ${#FILES[@]} -eq 0 ]]; then
  echo "Error: no files specified" >&2
  usage >&2
  exit 1
fi

for file in "${FILES[@]}"; do
  if [[ ! -f "$file" ]]; then
    echo "Warning: $file not found, skipping" >&2
    continue
  fi

  tmpfile=$(mktemp)
  trap "rm -f '$tmpfile'" EXIT

  awk '
  # Detect trigger line: conn = VERB(conn, ARGS)
  # where VERB is get/post/put/patch/delete/head/options
  /^[[:space:]]*conn = (get|post|put|patch|delete|head|options)\(conn, / {
    trigger_line = $0
    # Extract leading whitespace
    match($0, /^[[:space:]]*/)
    indent = substr($0, RSTART, RLENGTH)

    # Extract verb and args from: conn = verb(conn, args)
    rest = $0
    sub(/^[[:space:]]*conn = /, "", rest)
    # rest is now: verb(conn, args)
    paren_pos = index(rest, "(")
    verb = substr(rest, 1, paren_pos - 1)
    # args portion: everything after "conn, " up to the trailing ")"
    inner = substr(rest, paren_pos + 1)
    sub(/\)$/, "", inner)
    # inner is: conn, args
    sub(/^conn, /, "", inner)
    args = inner

    # Read the next non-blank line
    triggered = 1
    next
  }

  triggered == 1 {
    # Skip blank lines, accumulating them
    if ($0 ~ /^[[:space:]]*$/) {
      blank_lines = blank_lines $0 "\n"
      next
    }

    next_line = $0
    triggered = 0

    # Now look ahead: count how many subsequent lines (until scope boundary)
    # reference "conn" — to decide Case 4 vs Cases 1-3
    # We already have next_line. Check if next_line references conn.
    # Then peek further lines.

    # For simplicity: check if next_line matches Case 1, 2, or 3 patterns.
    # If it does, check the line AFTER that for more conn references (Case 4 override).

    # Case 1: var = helper(conn, status)
    # helpers: html_response, json_response, text_response, response, redirected_to
    case1 = 0
    if (match(next_line, /^[[:space:]]*([a-z_]+) = (html_response|json_response|text_response|response|redirected_to)\(conn, [^)]+\)$/, m1)) {
      case1 = 1
      c1_var = m1[1]
      c1_helper = m1[2]
      # Extract status from helper(conn, status)
      match(next_line, /\(conn, ([^)]+)\)/, m1s)
      c1_status = m1s[1]
    }

    # Case 2: assert helper(conn, status) with optional =~ "..."
    case2 = 0
    if (match(next_line, /^[[:space:]]*assert (html_response|json_response|text_response|response|redirected_to)\(conn, ([^)]+)\)(.*)$/, m2)) {
      case2 = 1
      c2_helper = m2[1]
      c2_status = m2[2]
      c2_tail = m2[3]
    }

    # Case 3: assert %{...} = helper(conn, status)
    case3 = 0
    if (match(next_line, /^[[:space:]]*assert (%\{[^}]*\}) = (html_response|json_response|text_response|response|redirected_to)\(conn, ([^)]+)\)$/, m3)) {
      case3 = 1
      c3_pattern = m3[1]
      c3_helper = m3[2]
      c3_status = m3[3]
    }

    # If we matched Case 1, 2, or 3, emit the merged line
    if (case1) {
      print indent c1_var " = conn |> " verb "(" args ") |> " c1_helper "(" c1_status ")"
      if (blank_lines != "") printf "%s", blank_lines
      blank_lines = ""
      next
    }
    if (case2) {
      print indent "assert conn |> " verb "(" args ") |> " c2_helper "(" c2_status ")" c2_tail
      if (blank_lines != "") printf "%s", blank_lines
      blank_lines = ""
      next
    }
    if (case3) {
      print indent "assert " c3_pattern " = conn |> " verb "(" args ") |> " c3_helper "(" c3_status ")"
      if (blank_lines != "") printf "%s", blank_lines
      blank_lines = ""
      next
    }

    # If next_line references conn at all, this is Case 4 territory
    # (multiple uses without a recognized single-merge pattern)
    if (next_line ~ /conn/) {
      # Case 4: rename to response
      print indent "response = conn |> " verb "(" args ")"
      if (blank_lines != "") printf "%s", blank_lines
      blank_lines = ""
      # Replace conn with response in next_line
      gsub(/conn/, "response", next_line)
      print next_line
      # Continue replacing conn->response in subsequent lines until scope boundary
      renaming = 1
      next
    }

    # No conn reference on next line — leave trigger unchanged (fallback)
    print trigger_line
    if (blank_lines != "") printf "%s", blank_lines
    blank_lines = ""
    print next_line
    next
  }

  # Renaming mode for Case 4: replace conn with response until scope boundary
  renaming == 1 {
    # Scope boundary: blank line, "end", reduced indentation, or new conn = assignment
    if ($0 ~ /^[[:space:]]*$/ || $0 ~ /^[[:space:]]*end$/ || $0 ~ /^[[:space:]]*conn =/) {
      renaming = 0
      print
      next
    }
    gsub(/conn/, "response")
    print
    next
  }

  # Normal mode: pass through
  {
    if (blank_lines != "") {
      printf "%s", blank_lines
      blank_lines = ""
    }
    print
  }

  BEGIN { triggered = 0; renaming = 0; blank_lines = "" }
  ' "$file" > "$tmpfile"

  if $DRY_RUN; then
    diff -u "$file" "$tmpfile" || true
  else
    mv "$tmpfile" "$file"
    echo "Refactored: $file"
  fi
done