#!/bin/bash

# AI-Fix: Code review and fixing tool
# Two-stage LLM workflow: review code, then apply fixes

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/ai-common"

show_usage() {
    cat << 'EOF'
Usage: ai-fix [-r] PROMPT_FILE [FILE...]
Review code according to a prompt and optionally apply fixes.

Arguments:
  PROMPT_FILE     File containing the review prompt/instructions

Options:
  -r, --review    Review only mode (don't apply fixes)
  -h              Show this help message

Examples:
  # Review and fix code (default behavior)
  ai-fix code-review-prompt.txt mycode.py

  # In vim, review and fix selected text
  :'<,'>!ai-fix code-review-prompt.txt

  # Review only (no fixes applied)
  ai-fix -r code-review-prompt.txt mycode.py
EOF
}

build_review_prompt() {
    local review_instructions="$1"
    
    echo "$review_instructions"
}

build_fix_prompt() {
    local review="$1"
    
    cat << EOF
You are a code fixing tool. Apply straightforward high-priority fixes to the code based on the provided review feedback.

Review Feedback:
$review
EOF
}

main() {
    local review_only=false
    local prompt_file=""
    local files=()
    local input review_instructions review fixed_code
    
    # Parse options
    while getopts "rh-:" opt; do
        case $opt in
            r) review_only=true ;;
            h) handle_help_option show_usage ;;
            -)
                case $OPTARG in
                    review) review_only=true ;;
                    *) handle_invalid_option "--$OPTARG" ;;
                esac
                ;;
            \?) handle_invalid_option "$OPTARG" ;;
        esac
    done
    
    shift $((OPTIND-1))
    
    # Validate environment and arguments
    ensure_dependencies
    ensure_argument_provided "Prompt file" "$1" show_usage
    
    prompt_file="$1"
    shift
    files=("$@")
    
    # Validate prompt file exists
    ensure_files_exist "$prompt_file"
    
    # Validate additional files if provided
    [[ ${#files[@]} -gt 0 ]] && ensure_files_exist "${files[@]}"
    
    # Read review instructions
    if ! review_instructions=$(cat "$prompt_file"); then
        print_error "Failed to read prompt file: $prompt_file"
        exit "$EXIT_USAGE_ERROR"
    fi
    
    # Process input code
    if ! input=$(process_input_sources "${files[@]}"); then
        exit "$EXIT_NO_MATCH"
    fi
    
    # Stage 1: Generate code review
    local review_prompt
    review_prompt=$(build_review_prompt "$review_instructions")
    
    if ! review=$(execute_llm_request "$review_prompt" "$input"); then
        case $? in
            "$EXIT_NO_MATCH") print_error "$(error_llm_no_response)" ;;
            "$EXIT_API_ERROR") print_error "$(error_llm_api_error)" ;;
            *) print_error "$(error_llm_command_failed)" ;;
        esac
        exit $?
    fi
    
    # If review-only mode, output the review and exit
    if $review_only; then
        echo "$review"
        exit "$EXIT_SUCCESS"
    fi
    
    # Stage 2: Apply fixes based on review
    local fix_prompt
    fix_prompt=$(build_fix_prompt "$review")
    
    if ! fixed_code=$(execute_llm_request "$fix_prompt" "$input"); then
        case $? in
            "$EXIT_NO_MATCH") print_error "$(error_llm_no_response)" ;;
            "$EXIT_API_ERROR") print_error "$(error_llm_api_error)" ;;
            *) print_error "$(error_llm_command_failed)" ;;
        esac
        exit $?
    fi
    
    # Output the fixed code
    echo "$fixed_code"
    exit "$EXIT_SUCCESS"
}

main "$@"