#!/bin/bash
# ============================================
# Task Complete Review Hook for Claude Code
#
# Triggers when a task is marked as completed via TodoWrite
# and sends the created/modified files to both Gemini MCP
# and Codex MCP for improvement suggestions.
#
# Hook Events Used:
#   - PostToolUse (TodoWrite): Detects task completions
#   - PostToolUse (Edit|Write): Tracks file modifications
#
# MCP Servers Used:
#   - mcp__gemini-cli__ask-gemini - Fast analysis with Gemini
#   - mcp__codex__codex - Deep analysis with Codex
#
# Created: 2024-12-10
# ============================================

set -uo pipefail
# Note: -e removed to allow graceful error handling

# Configuration
PROJECT_ROOT="/Users/Asim/Desktop/mawidi_codex"
TASK_QUEUE_DIR="/tmp/claude_task_review"
TASK_FILES_QUEUE="$TASK_QUEUE_DIR/task_files.jsonl"
CURRENT_TASK_FILE="$TASK_QUEUE_DIR/current_task.json"
LOG_FILE="$TASK_QUEUE_DIR/task-review.log"
MAX_FILE_SIZE=150000  # 150KB max per file for review
MAX_FILES_PER_REVIEW=10  # Limit files sent for review

# Ensure directories exist
mkdir -p "$TASK_QUEUE_DIR"

# Function to log messages
log_message() {
    local timestamp
    timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] $1" >> "$LOG_FILE"
}

# Function to check if file should be included in review
should_include_file() {
    local file_path="$1"

    # Only include code files
    if [[ ! "$file_path" =~ \.(ts|tsx|js|jsx|py|go|rs|java|sql|prisma|sh|mjs|cjs)$ ]]; then
        return 1
    fi

    # Skip directories that should not be reviewed
    local skip_dirs=("node_modules" ".next" ".git" "dist" "build" "coverage" "__pycache__" ".cache")
    for dir in "${skip_dirs[@]}"; do
        if [[ "$file_path" == *"/$dir/"* ]] || [[ "$file_path" == *"$dir/"* ]]; then
            return 1
        fi
    done

    # Skip very large files
    if [[ -f "$file_path" ]]; then
        local file_size
        file_size=$(stat -f%z "$file_path" 2>/dev/null || stat --printf="%s" "$file_path" 2>/dev/null || echo "0")
        if [[ "$file_size" -gt $MAX_FILE_SIZE ]]; then
            return 1
        fi
    fi

    return 0
}

# Function to get relative path
get_relative_path() {
    local file_path="$1"
    echo "${file_path#$PROJECT_ROOT/}"
}

# Function to read file content safely
read_file_content() {
    local file_path="$1"
    local max_lines=200  # Limit to first 200 lines for review

    if [[ -f "$file_path" ]]; then
        head -n "$max_lines" "$file_path" 2>/dev/null || cat "$file_path" 2>/dev/null || echo "[Unable to read file]"
    else
        echo "[File not found]"
    fi
}

# Function to handle file edits/writes - queue files associated with current task
handle_file_change() {
    local input="$1"

    local tool_name
    tool_name=$(echo "$input" | jq -r '.tool_name // empty')

    # Only handle Edit and Write tools
    if [[ "$tool_name" != "Edit" ]] && [[ "$tool_name" != "Write" ]]; then
        echo '{"continue": true}'
        exit 0
    fi

    local file_path
    file_path=$(echo "$input" | jq -r '.tool_input.file_path // empty')

    if [[ -z "$file_path" ]]; then
        echo '{"continue": true}'
        exit 0
    fi

    # Check if file should be included
    if ! should_include_file "$file_path"; then
        echo '{"continue": true}'
        exit 0
    fi

    log_message "File change: Queuing file for task review: $file_path"

    # Add to task files queue (compact single-line JSON for JSONL format)
    local queue_entry
    queue_entry=$(jq -cn \
        --arg timestamp "$(date -u +%Y-%m-%dT%H:%M:%SZ)" \
        --arg file_path "$file_path" \
        --arg tool_name "$tool_name" \
        '{timestamp: $timestamp, file_path: $file_path, tool_name: $tool_name}')

    echo "$queue_entry" >> "$TASK_FILES_QUEUE"

    echo '{"continue": true}'
}

# Function to handle TodoWrite - detect task completions
handle_todo_write() {
    local input="$1"

    local tool_name
    tool_name=$(echo "$input" | jq -r '.tool_name // empty' 2>/dev/null || echo "")

    if [[ "$tool_name" != "TodoWrite" ]]; then
        echo '{"continue": true}'
        exit 0
    fi

    log_message "TodoWrite detected, checking for completed tasks..."

    # Parse todos from tool_input - handle as string if needed
    local todos
    todos=$(echo "$input" | jq -c '.tool_input.todos // []' 2>/dev/null || echo "[]")

    # Find completed tasks
    local completed_tasks
    completed_tasks=$(echo "$todos" | jq -r '[.[] | select(.status == "completed")] | length' 2>/dev/null || echo "0")

    if [[ "$completed_tasks" -lt 1 ]] || [[ -z "$completed_tasks" ]]; then
        log_message "No completed tasks found"
        echo '{"continue": true}'
        exit 0
    fi

    # Get completed task descriptions
    local completed_descriptions
    completed_descriptions=$(echo "$todos" | jq -r '[.[] | select(.status == "completed") | .content] | join(", ")' 2>/dev/null || echo "Task completed")

    log_message "Completed tasks: $completed_descriptions"

    # Check if we have files queued for this task
    if [[ ! -f "$TASK_FILES_QUEUE" ]] || [[ ! -s "$TASK_FILES_QUEUE" ]]; then
        log_message "No files queued for review"
        echo '{"continue": true}'
        exit 0
    fi

    # Get unique files (limit to MAX_FILES_PER_REVIEW)
    local -a unique_files
    unique_files=()
    while IFS= read -r line; do
        local file_path
        file_path=$(echo "$line" | jq -r '.file_path // empty' 2>/dev/null || echo "")
        if [[ -n "$file_path" ]] && [[ -f "$file_path" ]]; then
            # Check if already in array
            local found=0
            for existing in "${unique_files[@]+"${unique_files[@]}"}"; do
                if [[ "$existing" == "$file_path" ]]; then
                    found=1
                    break
                fi
            done
            if [[ $found -eq 0 ]]; then
                unique_files+=("$file_path")
                if [[ ${#unique_files[@]} -ge $MAX_FILES_PER_REVIEW ]]; then
                    break
                fi
            fi
        fi
    done < "$TASK_FILES_QUEUE"

    local file_count=${#unique_files[@]}

    if [[ $file_count -lt 1 ]]; then
        log_message "No valid files to review after filtering"
        echo '{"continue": true}'
        exit 0
    fi

    log_message "Triggering AI review for $file_count files after task completion"

    # Build file list and content for review
    local file_list=""
    local file_contents=""
    local gemini_refs=""

    for file in "${unique_files[@]}"; do
        local rel_path
        rel_path=$(get_relative_path "$file")
        file_list="${file_list}- ${rel_path}\n"
        gemini_refs="${gemini_refs}@${rel_path} "

        # Get file content
        local content
        content=$(read_file_content "$file")
        file_contents="${file_contents}\n\n### File: ${rel_path}\n\`\`\`\n${content}\n\`\`\`"
    done

    # Create the system message that prompts for AI review
    local system_message
    system_message=$(cat <<EOF
## Task Completed - AI Code Review Recommended

**Completed Task(s):** ${completed_descriptions}

**Files Created/Modified (${file_count}):**
$(echo -e "$file_list")

### Recommended: Get AI Improvement Suggestions

To get comprehensive improvement suggestions, run these MCP tools:

**1. Gemini Analysis (Fast):**
\`\`\`
mcp__gemini-cli__ask-gemini with prompt:
"${gemini_refs}

Review these files and suggest improvements for:
1. Code quality and best practices
2. Performance optimizations
3. Security concerns
4. Error handling
5. Documentation needs

Provide specific, actionable suggestions with code examples."
\`\`\`

**2. Codex Deep Analysis:**
\`\`\`
mcp__codex__codex with prompt:
"Analyze and suggest improvements for these files:

${file_list}
Focus on:
- Architecture and design patterns
- Potential bugs or edge cases
- Maintainability improvements
- Type safety enhancements"
\`\`\`

Would you like me to run the AI code review on these completed files?
EOF
)

    # Clear the queue after processing
    rm -f "$TASK_FILES_QUEUE"

    # Return system message to prompt for review
    jq -n --arg msg "$system_message" '{"continue": true, "systemMessage": $msg}'
}

# Main handler
main() {
    # Read input from stdin
    local input
    input=$(cat)

    # Validate JSON input
    if ! echo "$input" | jq empty 2>/dev/null; then
        log_message "Invalid JSON input received"
        echo '{"continue": true}'
        exit 0
    fi

    # Get hook event name
    local hook_event
    hook_event=$(echo "$input" | jq -r '.hook_event_name // empty' 2>/dev/null || echo "")

    # Get tool name
    local tool_name
    tool_name=$(echo "$input" | jq -r '.tool_name // empty' 2>/dev/null || echo "")

    log_message "Hook triggered: $hook_event, Tool: $tool_name"

    case "$tool_name" in
        "TodoWrite")
            handle_todo_write "$input"
            ;;
        "Edit"|"Write")
            handle_file_change "$input"
            ;;
        *)
            echo '{"continue": true}'
            ;;
    esac
}

# Run main
main "$@"

exit 0
