#!/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 '
  function is_hard_scope_boundary(line) {
    return (line ~ /^[[:space:]]*end$/ || line ~ /^[[:space:]]*conn =/ || line ~ /^[[:space:]]*(test|describe) /)
  }

  function conn_used_ahead(start_idx,    i, line) {
    for (i = start_idx; i <= total_lines; i++) {
      line = lines[i]
      if (is_hard_scope_boundary(line)) return 0
      if (line ~ /conn/) return 1
    }
    return 0
  }

  # Read all lines into array
  { lines[NR] = $0 }

  END {
    total_lines = NR
    triggered = 0
    renaming = 0
    blank_lines = ""

    for (idx = 1; idx <= total_lines; idx++) {
      line = lines[idx]

      # If in renaming mode (Case 4 continuation)
      if (renaming) {
        if (is_hard_scope_boundary(line)) {
          renaming = 0
          # Fall through: do NOT continue — let the line be processed normally below
        } else {
          gsub(/conn/, "response", line)
          print line
          continue
        }
      }

      # If waiting for next non-blank line after trigger
      if (triggered) {
        if (line ~ /^[[:space:]]*$/) {
          blank_lines = blank_lines line "\n"
          continue
        }

        next_line = line
        triggered = 0

        # Case 1: var = helper(conn, status)
        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]
          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 Case 1/2/3 matched, check if conn is used further ahead
        # If so, fall through to Case 4
        if (case1 || case2 || case3) {
          if (conn_used_ahead(idx + 1)) {
            case1 = 0; case2 = 0; case3 = 0
          }
        }

        if (case1) {
          print indent c1_var " = conn |> " verb "(" args ") |> " c1_helper "(" c1_status ")"
          if (blank_lines != "") printf "%s", blank_lines
          blank_lines = ""
          continue
        }
        if (case2) {
          print indent "assert conn |> " verb "(" args ") |> " c2_helper "(" c2_status ")" c2_tail
          if (blank_lines != "") printf "%s", blank_lines
          blank_lines = ""
          continue
        }
        if (case3) {
          print indent "assert " c3_pattern " = conn |> " verb "(" args ") |> " c3_helper "(" c3_status ")"
          if (blank_lines != "") printf "%s", blank_lines
          blank_lines = ""
          continue
        }

        # If next_line references conn at all, this is Case 4 territory
        if (next_line ~ /conn/) {
          print indent "response = conn |> " verb "(" args ")"
          if (blank_lines != "") printf "%s", blank_lines
          blank_lines = ""
          gsub(/conn/, "response", next_line)
          print next_line
          renaming = 1
          continue
        }

        # 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
        continue
      }

      # Detect trigger line: conn = VERB(conn, ARGS)
      if (line ~ /^[[:space:]]*conn = (get|post|put|patch|delete|head|options)\(conn, /) {
        trigger_line = line
        match(line, /^[[:space:]]*/)
        indent = substr(line, RSTART, RLENGTH)

        rest = line
        sub(/^[[:space:]]*conn = /, "", rest)
        paren_pos = index(rest, "(")
        verb = substr(rest, 1, paren_pos - 1)
        inner = substr(rest, paren_pos + 1)
        sub(/\)$/, "", inner)
        sub(/^conn, /, "", inner)
        args = inner

        triggered = 1
        continue
      }

      # Normal mode: pass through
      if (blank_lines != "") {
        printf "%s", blank_lines
        blank_lines = ""
      }
      print line
    }
  }
  ' "$file" > "$tmpfile"

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