#!/usr/bin/env bash
# LICENSE: unlicense. This is free and unencumbered software released into the public domain.
# see unlicense.org for full license

set -euo pipefail

function usage {
  cat <<EOF
review [--verbose] [--context TEXT|-] [--help] [git-diff-arguments...]

Ask an LLM to review code changes. This tool passes arguments directly to 'git diff',
allowing you to use any git diff syntax or options.

Options:
  -c, --context TEXT  Add additional context for the review, appended to the system prompt. Will be concatenated if provided multiple times
      --context -     Read additional context from stdin
  -h, --help          Show this help message
  -v, --verbose       Enable verbose output

Review Examples:
  # Review unstaged changes
  review

  # Review with additional context
  review --context "Focus your review on possible authentication bypasses"

  # Review with context from stdin
  cat PR_DESCRIPTION.md | review --context -

  # Review with context from a command
  git log -1 --pretty=%B | review --context -

  # Review staged changes
  review --cached

  # Review changes between HEAD and main
  review main

  # Review changes between two branches
  review main feature-branch
  # OR
  review main..feature-branch

  # Review only changes since branch diverged from main
  review main...feature-branch

  # Review a remote branch
  review origin/main..origin/feature-branch

  # Limit review to specific files
  review main -- src/components/

  # Exclude a file (package-lock.json) from the review:
  review -- ':!package-lock.json'

  # Review from the merge base, and only *.js and *.ts files
  review next...HEAD -- '**/*.js' '**/*.ts'

  # Adjust context lines
  review -U5 main

Dot Notation:
  - Two dots (A..B): Direct comparison between A and B
  - Three dots (A...B): Compare common ancestor of A and B with B

Depends on:
- llm: https://github.com/simonw/llm
- bat: https://github.com/sharkdp/bat (optional)
EOF
  exit "${1:-0}"
}

git_args=()
has_unified_context=false
context_value=10
additional_context=""

# Process only our custom arguments, pass everything else to git
while [[ $# -gt 0 ]]; do
  case "$1" in
    -v|--verbose)
      set -x
      shift
      ;;
    -c|--context)
      shift
      if [[ $# -gt 0 ]]; then
        # Read from stdin for context
        if [[ "$1" == "-" ]]; then
          if [[ -t 0 ]]; then
            error "No stdin input available for --context -"
          fi
          new_context=$(cat)
          if [[ -z "$new_context" ]]; then
            error "Empty input from stdin for --context -"
          fi
        else
          new_context=$1
        fi
        if [[ -n "$additional_context" ]]; then
          additional_context="${additional_context}

$new_context"
        else
          additional_context="$new_context"
        fi
        shift
      else
        error "Missing value for --context option"
      fi
      ;;
    -U[0-9]*)
      # Form: -U10
      has_unified_context=true
      context_value="${1#-U}"
      git_args+=("$1")
      shift
      ;;
    -U)
      # Form: -U 10
      has_unified_context=true
      shift
      if [[ $# -gt 0 && "$1" =~ ^[0-9]+$ ]]; then
        context_value="$1"
        # normalize to `-U10` to ease our checking later on
        git_args+=("-U$1")
        shift
      else
        error "Missing value for -U option"
      fi
      ;;
    --unified=*)
      # Form: --unified=10
      has_unified_context=true
      context_value="${1#--unified=}"
      git_args+=("$1")
      shift
      ;;
    -h|--help)
      usage
      ;;
    *)
      # Store all other arguments to pass to git diff
      git_args+=("$1")
      shift
      ;;
  esac
done

readonly RED='\033[0;31m'
readonly BLUE='\033[0;34m'
readonly RESET='\033[0m'

info() {
  printf "${BLUE}• %s${RESET}\n" "$1" >&2
}

error() {
  printf "${RED}❌ %s${RESET}\n" "$1" >&2
  usage 1
}

if ! command -v llm >/dev/null 2>&1; then
  error "Missing required command llm. On mac: brew install llm"
fi

# Default unified context if none specified. The idea here is to increase the
# context (git defaults to 3 lines) so that the LLM has more context for its
# review. Later on we'll check if this generates too much output and shorten it
# if so
if [[ "$has_unified_context" == false ]]; then
  if [[ ${#git_args[@]} -eq 0 ]]; then
    git_args=("-U$context_value")
  else
    git_args=("-U$context_value" "${git_args[@]}")
  fi
fi

# create a tempfile for storing git error output
git_error=$(mktemp)
trap 'rm -f "$git_error"' EXIT

# Run git diff and append any additional context to it
if ! diff_output=$(git diff "${git_args[@]}" 2>"$git_error"); then
  printf "${RED}❌ Git diff command failed: git diff %s${RESET}\n" "${git_args[*]}" >&2
  cat "$git_error" >&2
  exit 1
fi
diff_output="${diff_output}

${additional_context}"

if [[ -z "$diff_output" ]]; then
  error "No changes found to review."
fi

# I wish there were a simple consistent method to count tokens, but there isn't
# as far as I can tell, so we're gonna use a poor estimation and keep safely
# inside the context limit
max_tokens=50000  # Claude's limit is 100k, this should be a safe amount
chars_per_token=4 # simple approximation

# Estimate token count and reduce context if needed
char_count=${#diff_output}
estimated_tokens=$((char_count / chars_per_token))

if [[ $estimated_tokens -gt $max_tokens ]]; then
  # Calculate reduced context
  reduced_context=$((context_value * max_tokens / estimated_tokens))
  reduced_context=$((reduced_context > 0 ? reduced_context : 1))
  
  info "Reducing context to $reduced_context lines to fit token limits"
  
  # Replace unified context in git args
  new_git_args=()
  for arg in "${git_args[@]}"; do
    if [[ "$arg" =~ ^-U[0-9]+$ ]]; then
      new_git_args+=("-U$reduced_context")
    elif [[ "$arg" =~ ^--unified=[0-9]+$ ]]; then
      new_git_args+=("--unified=$reduced_context")
    else
      new_git_args+=("$arg")
    fi
  done

  if [[ $((${#diff_output} / chars_per_token)) -gt max_tokens ]]; then
    error "Diff is too large to process even with minimal context. Try reviewing a smaller set of changes."
  fi
  
  # Re-run git diff with reduced context
  if ! diff_output=$(git diff "${new_git_args[@]}" 2>"$git_error"); then
    printf "${RED}❌ Git diff command failed with reduced context: git diff %s${RESET}\n" "${new_git_args[*]}" >&2
    cat "$git_error" >&2
    exit 1
  fi
fi

prompt="Please review this PR as if you were a senior engineer.

## Focus Areas
- Architecture and design decisions
- Potential bugs and edge cases
- Performance considerations
- Security implications
- Code maintainability and best practices
- Test coverage

## Review Format
- Start with a brief summary of the PR purpose and changes
- List strengths of the implementation
- Identify issues and improvement opportunities (ordered by priority)
- Provide specific code examples for suggested changes where applicable

Please be specific, constructive, and actionable in your feedback. Output the review in markdown format."

if command -v bat >/dev/null 2>&1; then
  echo "$diff_output" | llm -s "$prompt" | bat --paging=never --style=plain --language=markdown
else
  echo "$diff_output" | llm -s "$prompt"
fi
exit_status=${PIPESTATUS[1]}

if [[ "$exit_status" -eq 130 ]]; then
  # User pressed Ctrl+C, exit silently
  exit 130
elif [[ "$exit_status" -ne 0 ]]; then
  echo "Error: LLM command failed." >&2
  exit 1
fi


