#!/bin/bash

# Peer Review System
# This script implements a peer review process where one model provides an initial response
# and other models review and suggest improvements through iterative refinement.
#
# APPLICATION LOGIC:
# The peer review process implements a collaborative refinement system where AI models
# provide structured feedback to improve response quality. The system operates through
# three distinct phases designed to enhance clarity, accuracy, and completeness:
#
# PHASE 1 - INITIAL RESPONSE GENERATION:
#   - A randomly selected model generates the first response to the user's prompt
#   - Random selection helps prevent bias toward any particular model
#   - The initial response serves as the foundation for peer review
#   - This creates a starting point that can be improved through collective feedback
#
# PHASE 2 - PEER REVIEW:
#   - Other models analyze the initial response and provide structured feedback
#   - Reviewers suggest specific edits, clarifications, and improvements
#   - Feedback focuses on clarity, completeness, accuracy, and organization
#   - Multiple perspectives may identify different areas for improvement
#
# PHASE 3 - RESPONSE REFINEMENT:
#   - The original responding model incorporates peer feedback to create an improved response
#   - The model may revise, expand, clarify, or reorganize based on suggestions
#   - Iterative improvement may address multiple rounds of feedback
#   - The process emphasizes collaborative enhancement rather than replacement
#
# PEER REVIEW MODELING:
# The system applies academic peer review principles to AI response refinement:
#   - Random author selection helps prevent systematic bias in initial responses
#   - Multiple reviewers provide diverse perspectives on the same work
#   - Structured feedback focuses on specific improvements rather than general criticism
#   - Author retains control over final revisions while considering peer input
#   - Transparency through logging shows the evolution of the response
#   - The process may help catch errors, improve clarity, or enhance completeness
#
# The peer review process continues for a configurable number of iterations,
# with each cycle potentially improving upon the previous version.
# The system emphasizes collaborative improvement through structured feedback and revision.

# --- Model Configuration ---
MODELS=(
    "llama3:8b-instruct-q4_K_M"
    "phi3:3.8b-mini-4k-instruct-q4_K_M"
    "deepseek-r1:1.5b"
    "gemma3n:e2b"
    "dolphin3:latest"
)

# --- Defaults ---
DEFAULT_ITERATIONS=1

# --- Argument Validation ---
if [ "$#" -lt 1 ]; then
    echo -e "\n\tPeer Review"
    echo -e "\tThis script implements a peer review process where one model provides an initial response"
    echo -e "\tand other models review and suggest improvements through iterative refinement."
    echo -e "\n\tUsage: $0 [-f <file_path>] \"<your prompt>\" [number_of_review_iterations]"
    echo -e "\n\tExample: $0 -f ./input.txt \"Please analyze this text\" 1"
    echo -e "\n\tIf number_of_review_iterations is not provided, the program will default to $DEFAULT_ITERATIONS iterations."
    echo -e "\n\t-f <file_path> (optional): Append the contents of the file to the prompt."
    echo -e "\n"
    exit 1
fi

# --- Argument Parsing ---
FILE_PATH=""
while getopts "f:" opt; do
  case $opt in
    f)
      FILE_PATH="$OPTARG"
      ;;
    *)
      echo "Invalid option: -$OPTARG" >&2
      exit 1
      ;;
  esac
done
shift $((OPTIND -1))

PROMPT="$1"
if [ -z "$2" ]; then
    ITERATIONS=$DEFAULT_ITERATIONS
else
    ITERATIONS=$2
fi

# If file path is provided, append its contents to the prompt
if [ -n "$FILE_PATH" ]; then
    if [ ! -f "$FILE_PATH" ]; then
        echo "File not found: $FILE_PATH" >&2
        exit 1
    fi
    FILE_CONTENTS=$(cat "$FILE_PATH")
    PROMPT="$PROMPT\n[FILE CONTENTS]\n$FILE_CONTENTS\n[END FILE]"
fi

# --- File Initialization ---
# Create a temporary directory if it doesn't exist
mkdir -p ~/tmp
# Create a unique file for this session based on the timestamp
SESSION_FILE=~/tmp/peer-review_$(date +%Y%m%d_%H%M%S).txt

echo "Peer Review Session Log: ${SESSION_FILE}"
echo "---------------------------------"

# Store the initial user prompt in the session file
echo "USER PROMPT: ${PROMPT}" >> "${SESSION_FILE}"
echo "" >> "${SESSION_FILE}"
echo "Processing peer review with ${ITERATIONS} iteration(s)..."

# --- Random Author Selection ---
AUTHOR_INDEX=$((RANDOM % ${#MODELS[@]}))
AUTHOR_MODEL="${MODELS[$AUTHOR_INDEX]}"

echo "Author model selected: ${AUTHOR_MODEL}"
echo "AUTHOR MODEL: ${AUTHOR_MODEL}" >> "${SESSION_FILE}"
echo "" >> "${SESSION_FILE}"

# --- Initial Response Generation ---
echo "Generating initial response..."
echo "INITIAL RESPONSE GENERATION:" >> "${SESSION_FILE}"
echo "============================" >> "${SESSION_FILE}"

INITIAL_PROMPT="You are an expert assistant. Please provide a comprehensive response to the following prompt. Be thorough, clear, and well-organized in your response.

PROMPT: ${PROMPT}"

INITIAL_RESPONSE=$(ollama run "${AUTHOR_MODEL}" "${INITIAL_PROMPT}")

echo "INITIAL RESPONSE (${AUTHOR_MODEL}):" >> "${SESSION_FILE}"
echo "${INITIAL_RESPONSE}" >> "${SESSION_FILE}"
echo "" >> "${SESSION_FILE}"

# --- Peer Review Iterations ---
CURRENT_RESPONSE="${INITIAL_RESPONSE}"

for iteration in $(seq 1 "${ITERATIONS}"); do
    echo "Starting peer review iteration ${iteration} of ${ITERATIONS}..."
    echo "PEER REVIEW ITERATION ${iteration}:" >> "${SESSION_FILE}"
    echo "=================================" >> "${SESSION_FILE}"
    
    # --- Step 1: Generate Peer Reviews ---
    echo "Step 1: Generating peer reviews..."
    echo "STEP 1 - PEER REVIEWS:" >> "${SESSION_FILE}"
    
    declare -a reviews
    declare -a reviewer_names
    
    review_count=0
    
    for i in "${!MODELS[@]}"; do
        # Skip the author model
        if [ "$i" -eq "$AUTHOR_INDEX" ]; then
            continue
        fi
        
        model="${MODELS[$i]}"
        echo "  Getting peer review from ${model}..."
        
        REVIEW_PROMPT="You are a peer reviewer. Your task is to provide constructive feedback on the following response. Focus on specific suggestions for improvement.

REVIEW GUIDELINES:
- Suggest specific edits or clarifications
- Identify areas that could be made clearer or more concise
- Point out any inaccuracies or missing information
- Suggest improvements to organization or structure
- Be constructive and specific in your feedback

RESPONSE TO REVIEW: ${CURRENT_RESPONSE}

Please provide your peer review feedback in a clear, structured format. Focus on actionable suggestions for improvement."

        review_output=$(ollama run "${model}" "${REVIEW_PROMPT}")
        
        reviews[$review_count]="${review_output}"
        reviewer_names[$review_count]="${model}"
        
        echo "REVIEW ${review_count+1} (${model}):" >> "${SESSION_FILE}"
        echo "${review_output}" >> "${SESSION_FILE}"
        echo "" >> "${SESSION_FILE}"
        
        review_count=$((review_count + 1))
    done
    
    # --- Step 2: Generate Refined Response ---
    echo "Step 2: Generating refined response based on peer feedback..."
    echo "STEP 2 - RESPONSE REFINEMENT:" >> "${SESSION_FILE}"
    
    # Combine all reviews for the author
    COMBINED_REVIEWS=""
    for i in $(seq 0 $((review_count - 1))); do
        COMBINED_REVIEWS="${COMBINED_REVIEWS}

REVIEW FROM ${reviewer_names[$i]}:
${reviews[$i]}"
    done
    
    REFINE_PROMPT="You are the author of the following response. Your peers have provided constructive feedback to help improve your work. Please revise your response based on their suggestions.

ORIGINAL PROMPT: ${PROMPT}
YOUR CURRENT RESPONSE: ${CURRENT_RESPONSE}
PEER REVIEW FEEDBACK: ${COMBINED_REVIEWS}

Please provide a revised version of your response that:
- Incorporates the constructive feedback from your peers
- Addresses specific suggestions for improvement
- Maintains your original insights while enhancing clarity and completeness
- Shows how you've responded to the peer review process"

    REFINED_RESPONSE=$(ollama run "${AUTHOR_MODEL}" "${REFINE_PROMPT}")
    
    echo "REFINED RESPONSE (${AUTHOR_MODEL}):" >> "${SESSION_FILE}"
    echo "${REFINED_RESPONSE}" >> "${SESSION_FILE}"
    echo "" >> "${SESSION_FILE}"
    
    # Update the current response for the next iteration
    CURRENT_RESPONSE="${REFINED_RESPONSE}"
    
    echo "Peer review iteration ${iteration} complete."
    echo "" >> "${SESSION_FILE}"
done

# --- Final Summary Generation ---
echo "Generating final summary..."
echo "FINAL SUMMARY GENERATION:" >> "${SESSION_FILE}"
echo "========================" >> "${SESSION_FILE}"

SUMMARY_PROMPT="You are an expert analyst. Based on the peer review process below, please provide a concise summary of the key improvements made and the overall quality of the final response.

ORIGINAL PROMPT: ${PROMPT}
FINAL REFINED RESPONSE: ${CURRENT_RESPONSE}

Please provide a summary that:
- Highlights the most significant improvements made through peer review
- Notes the quality and effectiveness of the final response
- Captures the collaborative nature of the peer review process
- Is clear, concise, and well-organized"

FINAL_SUMMARY=$(ollama run "${AUTHOR_MODEL}" "${SUMMARY_PROMPT}")

echo "FINAL SUMMARY (${AUTHOR_MODEL}):" >> "${SESSION_FILE}"
echo "${FINAL_SUMMARY}" >> "${SESSION_FILE}"

# --- Final Output ---
echo "---------------------------------"
echo "Peer review process complete."
echo "Final response:"
echo "---------------------------------"

echo "${CURRENT_RESPONSE}"
echo ""
echo "Author: ${AUTHOR_MODEL}"
echo "Peer Review Summary:"
echo "${FINAL_SUMMARY}"
echo ""
echo "Full peer review log: ${SESSION_FILE}"