#!/bin/bash

# Critique System
# This script uses a sequence of LLM calls to refine an initial response through critique and revision.
#
# APPLICATION LOGIC:
# The critique process implements an iterative refinement system where AI models
# collaborate to improve response quality through critique and revision. The system
# operates through three distinct phases designed to enhance clarity and accuracy:
#
# PHASE 1 - INITIAL RESPONSE GENERATION:
#   - A response model generates the first answer to the user's prompt
#   - The model is instructed to be honest about knowledge limitations
#   - This creates a baseline response that can be improved through iteration
#   - The initial response serves as the foundation for refinement
#
# PHASE 2 - CRITICAL REVIEW:
#   - A critic model analyzes the current response for potential issues
#   - The critic identifies misunderstandings, unclear areas, and improvement opportunities
#   - Constructive feedback focuses on specific problems rather than general criticism
#   - The critique provides targeted guidance for the refinement phase
#
# PHASE 3 - RESPONSE REFINEMENT:
#   - A refine model incorporates the critique to generate an improved response
#   - The refine model considers both the original prompt and the feedback
#   - Iterative improvement may address clarity, accuracy, or completeness issues
#   - Multiple refinement loops may progressively enhance response quality
#
# REFINEMENT MODELING:
# The system applies iterative improvement principles to AI response generation:
#   - Separate models for different roles may provide specialized perspectives
#   - Critical review helps identify blind spots in the initial response
#   - Iterative refinement allows for progressive improvement over multiple cycles
#   - Transparency through logging shows the evolution of the response
#   - The process may help catch errors or improve clarity that single-pass generation misses
#
# The refinement process continues for a configurable number of loops,
# with each iteration potentially improving upon the previous response.
# The system emphasizes quality improvement through structured feedback and revision.

# --- Model Configuration ---
RESPONSE_MODEL="llama3:8b-instruct-q4_K_M"
CRITIC_MODEL="phi3:3.8b-mini-4k-instruct-q4_K_M"
REFINE_MODEL="llama3:8b-instruct-q4_K_M"

# --- Defaults ---
DEFAULT_LOOPS=2

# --- Argument Validation ---
if [ "$#" -lt 1 ]; then
    echo -e "\n\tCritique"
    echo -e "\tThis script uses a sequence of LLM calls to refine an initial response through critique and revision."
    echo -e "\n\tUsage: $0 [-f <file_path>] \"<your prompt>\" [number_of_refinement_loops]"
    echo -e "\n\tExample: $0 -f ./input.txt \"Please summarize this text file\" 2"
    echo -e "\n\tIf number_of_refinement_loops is not provided, the program will default to $DEFAULT_LOOPS loops."
    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
    LOOPS=$DEFAULT_LOOPS
else
    LOOPS=$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/critique_$(date +%Y%m%d_%H%M%S).txt

echo "Critique Session Log: ${SESSION_FILE}"
echo "---------------------------------"

# --- Initial Prompt & Response ---

# 1. Store the initial user prompt in the session file
echo "USER PROMPT: ${PROMPT}" >> "${SESSION_FILE}"
echo "" >> "${SESSION_FILE}" # Add a newline for readability
echo "Processing initial response..."

# 2. The RESPONSE model generates the first answer
RESPONSE_PROMPT="You are an expert, curious assistant who isn't afraid to say when they don't know something. Please respond directly to the following prompt: ${PROMPT}"
RESPONSE_OUTPUT=$(ollama run "${RESPONSE_MODEL}" "${RESPONSE_PROMPT}")

# Append the response to the session file
echo "INITIAL RESPONSE (${RESPONSE_MODEL}):" >> "${SESSION_FILE}"
echo "${RESPONSE_OUTPUT}" >> "${SESSION_FILE}"
echo "" >> "${SESSION_FILE}"

# --- Refinement Loop ---

# This variable will hold the most recent response for the next loop iteration
CURRENT_RESPONSE="${RESPONSE_OUTPUT}"

for i in $(seq 1 "${LOOPS}"); do
    echo "Starting refinement loop ${i} of ${LOOPS}..."

    # 3. The CRITIC model reviews the last response
    CRITIC_PROMPT="You are a detail oriented, close reading, keenly critical reviewer. Your task is to raise questions, flag potential misunderstandings, and areas for improved clarity in the following text. Provide concise, constructive criticism. Do not rewrite the text, only critique it. TEXT TO CRITIQUE: ${CURRENT_RESPONSE}"
    CRITIC_OUTPUT=$(ollama run "${CRITIC_MODEL}" "${CRITIC_PROMPT}")

    # Append the critique to the session file
    echo "CRITICISM ${i} (${CRITIC_MODEL}):" >> "${SESSION_FILE}"
    echo "${CRITIC_OUTPUT}" >> "${SESSION_FILE}"
    echo "" >> "${SESSION_FILE}"

    # 4. The REFINE model reads the original prompt and the critique to generate a new response
    REFINE_PROMPT="You are an expert assistant. Your previous response was reviewed and critiqued. Your task now is to generate a refined, improved response to the original prompt based on the feedback provided. ORIGINAL PROMPT: ${PROMPT} CONSTRUCTIVE CRITICISM: ${CRITIC_OUTPUT} Generate the refined response now."
    REFINE_OUTPUT=$(ollama run "${REFINE_MODEL}" "${REFINE_PROMPT}")

    # Append the refined response to the session file
    echo "REFINED RESPONSE ${i} (${REFINE_MODEL}):" >> "${SESSION_FILE}"
    echo "${REFINE_OUTPUT}" >> "${SESSION_FILE}"
    echo "" >> "${SESSION_FILE}"

    # Update the current response for the next loop or for the final output
    CURRENT_RESPONSE="${REFINE_OUTPUT}"
done

# --- Final Output ---

echo "---------------------------------"
echo "Critique process complete."
echo "Final refined answer:"
echo "---------------------------------"

# Print the final, most refined answer to standard output
echo "${CURRENT_RESPONSE}"