#!/bin/bash

# --- CONFIGURATION ---
CODE_DIR=$(pwd)
PLAN_DIR="/mnt/ssd_data/agent_data" # This directory must exist
TASK_FILE="$CODE_DIR/current_task.txt"
# Create the necessary directory before attempting to write logs
mkdir -p "$PLAN_DIR" || { echo "FATAL ERROR: Could not create directory $PLAN_DIR. Check permissions." >&2; exit 1; }
LOG_FILE="$PLAN_DIR/agent_run_$(date +%Y%m%d_%H%M%S).log" # Unique log file per attempt

# --- CHECK FOR TASK ARGUMENT ---
if [ -z "$1" ]; then
    echo "Usage: $0 \"<Your development task description>\""
    exit 1
fi

# Initial task is the command-line argument
TASK="$1"

# Define files for the AI context once
CONTEXT_FILES="$CODE_DIR/bugs.md @$CODE_DIR/ChangeLog.md @$CODE_DIR/roadmap.md @$CODE_DIR/gemini.md"

# --- CORE AGENT WORKFLOW FUNCTION ---
run_agent_workflow() {
    local TASK_PROMPT="$1"
    
    echo "--- 🤖 Starting New Agent Attempt ---" | tee -a "$LOG_FILE"
    echo "Task: $TASK_PROMPT" | tee -a "$LOG_FILE"
    
    # 1. Human-in-the-Loop: Pre-flight Backup
    echo -e "\n--- 💾 ACTION REQUIRED: SAVE PROJECT STATE ---" | tee -a "$LOG_FILE"
    echo "Please run your custom version control script NOW to save the current state."
    read -p "Press ENTER when the state has been successfully saved..."
    echo "Backup confirmed. Proceeding with AI planning..." | tee -a "$LOG_FILE"

    PLAN_FILE="$PLAN_DIR/plan_$(date +%Y%m%d_%H%M%S).md"
    touch "$PLAN_FILE"

    # --- AGENT 1: PLANNER ---
    echo -e "\n--- 🧠 AGENT 1: PLANNER ---" | tee -a "$LOG_FILE"
    
    # Create the Planner Query
    PLANNER_QUERY=$(cat <<END_OF_PLANNER_PROMPT
Your role is the Planner. Based on the overall task: '$TASK_PROMPT', the project code (read the current directory content), and the content of all provided files, create a numbered, detailed development plan. Save ONLY the plan content. Do not include any conversational text or summary.
END_OF_PLANNER_PROMPT
    )

    # Execute Planner: Pass the query via echo to stdin and file context via '@'.
    PLANNER_OUTPUT=$(echo "$PLANNER_QUERY" | gemini $CONTEXT_FILES 2>&1)
    
    # The response is saved to the plan file (overwriting with only the AI's response)
    echo "$PLANNER_OUTPUT" > "$PLAN_FILE"
    echo "$PLANNER_OUTPUT" >> "$LOG_FILE" # Save full output to log
    
    # Check if the plan file is empty again
    if [ ! -s "$PLAN_FILE" ]; then
        echo "ERROR: Planner failed to create a non-empty plan file. Aborting." | tee -a "$LOG_FILE"
        return 2 # Exit code for failed plan
    fi

    # --- AGENT 2: EXECUTER (Loop through plan steps) ---
    echo -e "\n--- 💻 AGENT 2: EXECUTER (Working on Plan) ---" | tee -a "$LOG_FILE"
    
    # Read plan steps
    mapfile -t PLAN_STEPS < <(grep -E '^[0-9]+\.' "$PLAN_FILE" | sed 's/^[[:space:]]*//')
    STEP_COUNT=${#PLAN_STEPS[@]}
    
    if [ "$STEP_COUNT" -eq 0 ]; then
        echo "ERROR: Plan file is empty or unreadable. Aborting execution." | tee -a "$LOG_FILE"
        return 3 # Exit code for bad plan
    fi

    for i in "${!PLAN_STEPS[@]}"; do
        STEP_NUM=$((i + 1))
        STEP_CONTENT="${PLAN_STEPS[$i]}"
        echo "Executing Step $STEP_NUM/$STEP_COUNT: $STEP_CONTENT" | tee -a "$LOG_FILE"
        
        # Create Executer Query
        EXECUTER_QUERY=$(cat <<END_OF_EXECUTER_PROMPT
Your role is the Executer. The overall task is '$TASK_PROMPT'. The complete plan is in $PLAN_FILE. Execute ONLY step $STEP_NUM: '$STEP_CONTENT'. Make all necessary code changes in the $CODE_DIR directory. After executing the step, append a note starting with 'EXECUTER NOTE:' about what you did to the end of $PLAN_FILE. Do not output anything else.
END_OF_EXECUTER_PROMPT
        )
        
        # Execute Executer: Pass query to stdin, file context, and the plan file for notes.
        # NOTE: Using the -y flag (Yolo mode) is often necessary for non-interactive code modification with gemini.
        EXECUTER_OUTPUT=$(echo "$EXECUTER_QUERY" | gemini -y $CONTEXT_FILES @$PLAN_FILE 2>&1)
        echo "$EXECUTER_OUTPUT" >> "$LOG_FILE"

        # The executer output *itself* is discarded/logged, as the note is appended to $PLAN_FILE by the AI.
    done

    # --- AGENT 3: REVIEWER/DOCUMENTER ---
    echo -e "\n--- ✨ AGENT 3: REVIEWER & DOCUMENTER ---" | tee -a "$LOG_FILE"

    # Create Reviewer Query
    REVIEWER_QUERY=$(cat <<END_OF_REVIEWER_PROMPT
Your role is the Reviewer and Documenter. The original task was '$TASK_PROMPT'. The development is finished.
1. Score: On the first line ONLY, write the score in the format 'SCORE: [number]/10'.
2. Review: Provide feedback on code quality and completion.
3. Document: Update the necessary sections in $CODE_DIR/bugs.md, $CODE_DIR/ChangeLog.md, and $CODE_DIR/roadmap.md directly.
4. Re-prompt (Conditional): If the score is less than 7.5, suggest a better-worded prompt for the next attempt. Start this suggestion with 'NEW_PROMPT:'.
END_OF_REVIEWER_PROMPT
    )

    # Execute Reviewer
    # Use -y for documentation edits. Pass all files. Capture output for score parsing.
    REVIEW_OUTPUT=$(echo "$REVIEWER_QUERY" | gemini -y $CONTEXT_FILES @$PLAN_FILE 2>&1)
    echo "$REVIEW_OUTPUT" >> "$LOG_FILE"

    # --- SCORE PARSING AND DECISION ---
    # Extract the score line anywhere in the output
    # Use 'grep -E' to find the line containing "SCORE: [number]/10"
    SCORE_LINE=$(echo "$REVIEW_OUTPUT" | grep -E 'SCORE: [0-9]+\.?[0-9]*\/10')
    
    # Extract only the number preceding '/10' from that line
    SCORE=$(echo "$SCORE_LINE" | sed -n 's/.*SCORE: \([0-9]*\.[0-9]\|[0-9]\)\/10.*/\1/p')
    
    if [ -z "$SCORE" ]; then
        echo "ERROR: Could not parse score from reviewer output. Treating as failure (0.0)." | tee -a "$LOG_FILE"
        SCORE=0.0
    fi
    
    # Convert to float for comparison
    FLOAT_SCORE=$(echo "$SCORE * 10" | bc)
    THRESHOLD=75 # 7.5 * 10
    
    echo -e "\n--- DECISION: Final Score is $SCORE/10 ---" | tee -a "$LOG_FILE"

    if (( $(echo "$FLOAT_SCORE >= $THRESHOLD" | bc -l) )); then
        echo "✅ SUCCESS! Score meets the 7.5/10 threshold." | tee -a "$LOG_FILE"
        echo "Please run your custom version control script NOW to save the final successful state."
        read -p "Press ENTER when the final save is complete..."
        return 0 # Success
    else
        echo "❌ FAILURE! Score is below 7.5/10. Reverting and retrying..." | tee -a "$LOG_FILE"
        
        # Human-in-the-Loop: Revert State
        echo -e "\n--- 🚨 ACTION REQUIRED: REVERT PROJECT STATE ---" | tee -a "$LOG_FILE"
        echo "Please run your custom version control script NOW to revert to the previous saved state."
        read -p "Press ENTER when the revert is complete..."
        echo "Revert confirmed. Retrying..." | tee -a "$LOG_FILE"
        
        # Find New Prompt
        NEW_PROMPT=$(echo "$REVIEW_OUTPUT" | grep -m 1 "NEW_PROMPT:" | sed 's/.*NEW_PROMPT:[[:space:]]*//')
        
        if [ -z "$NEW_PROMPT" ]; then
            echo "ERROR: Reviewer did not suggest a new prompt. Using original task for retry." | tee -a "$LOG_FILE"
            TASK_TO_RETRY="$TASK"
        else
            echo "Using AI-suggested new prompt for retry: '$NEW_PROMPT'" | tee -a "$LOG_FILE"
            TASK_TO_RETRY="$NEW_PROMPT"
        fi
        
        return 1 # Failure, but proceed to retry loop
    fi
}

# --- MAIN RETRY LOOP ---
while true; do
    # Use the current TASK variable (either original or updated)
    run_agent_workflow "$TASK"
    EXIT_CODE=$?
    
    if [ "$EXIT_CODE" -eq 0 ]; then
        echo "Workflow finished successfully."
        break # Success
    elif [ "$EXIT_CODE" -eq 2 ] || [ "$EXIT_CODE" -eq 3 ]; then
        echo "Workflow aborted due to fatal error (see log $LOG_FILE)."
        break # Fatal error, stop
    else # EXIT_CODE is 1 (Failed score), so retry with new prompt
        # The 'run_agent_workflow' function has already set the 'TASK_TO_RETRY' variable.
        TASK="$TASK_TO_RETRY"
        echo -e "\n--- RETRYING WORKFLOW with new task: '$TASK' ---\n"
        sleep 5 # Pause before retry
    fi
done