#!/bin/sh
set -e

# --- Usage ---
if [ "$1" = "--usage" ]; then
  echo "Regenerates a past response (default: last) based on new instructions."
  exit 0
fi

# Ensure jq is installed
if ! command -v jq > /dev/null 2>&1; then
    echo "Error: 'jq' is required for this addon."
    exit 1
fi

# Default index
INDEX="-1"

# If the first arg is a number, treat it as the index
case "$1" in
    ''|*[!0-9-]*) ;;
    *) INDEX="$1"; shift ;;
esac

MODEL_FLAG=""
INSTRUCTION=""

while [ $# -gt 0 ]; do
    key="$1"
    case $key in
        --model)
        MODEL_FLAG="--model $2"
        shift
        shift
        ;;
        *)
        if [ -z "$INSTRUCTION" ]; then
            INSTRUCTION="$1"
        else
            INSTRUCTION="$INSTRUCTION $1"
        fi
        shift
        ;;
    esac
done

# Validate that we have an instruction
if [ -z "$INSTRUCTION" ]; then
  echo "Error: Missing instruction."
  echo "Usage: aico refine [index] \"<instruction>\" [--model <model>]"
  exit 1
fi

# 1. Create a temp file for the JSON data
JSON_DUMP=$(mktemp -t aico_refine_dump.XXXXXX)

# 2. Direct output to file instead of variable
if ! aico last "$INDEX" --json > "$JSON_DUMP" 2>/dev/null; then
    echo "Error: Could not retrieve message pair at index $INDEX."
    rm -f "$JSON_DUMP"
    exit 1
fi

# 3. Validation: Check if file is empty
if [ ! -s "$JSON_DUMP" ]; then
    echo "Error: Could not retrieve message pair at index $INDEX (Empty response)."
    rm -f "$JSON_DUMP"
    exit 1
fi

# 4. Read directly from file (No 'echo' pipe)
ORIG_MODE=$(jq -r .assistant.mode "$JSON_DUMP")
ORIG_CONTENT=$(jq -r .assistant.content "$JSON_DUMP")
ORIG_PROMPT_CONTENT=$(jq -r .user.content "$JSON_DUMP")
ABS_INDEX=$(jq -r .pair_index "$JSON_DUMP")
ORIG_USER_ID=$(jq -r .user.id "$JSON_DUMP")

# Ephemeral Execution Setup
SESSION_NAME="refine-$(date +%s)-$$"
PAYLOAD=$(mktemp -t aico_refine_payload.XXXXXX)

# Cleanup the dump file early
rm -f "$JSON_DUMP"
PROMPT_FILE=$(mktemp -t aico_refine_prompt.XXXXXX)
SYSTEM_PROMPT_FILE=$(mktemp -t aico_refine_sysprompt.XXXXXX)

cleanup() {
    rm -f "$PAYLOAD" "$PROMPT_FILE" "$SYSTEM_PROMPT_FILE"
}

# Define System Prompt based on mode and write to temp file
if [ "$ORIG_MODE" = "diff" ]; then
    cat <<EOF > "$SYSTEM_PROMPT_FILE"
You are a Code Refinement Engine.
Your goal is to regenerate a code patch based on new instructions.

**INPUT DATA:**
- <original_prompt>: The user's original request.
- <previous_draft>: The code you generated previously (which was incorrect or incomplete).
- <critique>: The user's specific instructions on how to fix the draft.

**CRITICAL RULES:**
1. You must output **ONLY** the valid \`SEARCH/REPLACE\` blocks.
2. Do **NOT** include conversational filler (e.g., 'Here is the fixed code').
3. The <critique> **OVERRIDES** any conflicting instructions in the <original_prompt>.
4. Your output will be spliced directly into the history as if it were your first and only response.
EOF
else
    cat <<EOF > "$SYSTEM_PROMPT_FILE"
You are a Content Refinement Engine.
Your goal is to rewrite a previous response to satisfy a critique.

**INPUT DATA:**
- <original_prompt>: The user's original request.
- <previous_draft>: The text you generated previously.
- <critique>: The user's specific instructions on how to change the draft.

**CRITICAL RULES:**
1. You must output **ONLY** the new response content.
2. Do **NOT** include conversational filler (e.g., 'I cannot edit history').
3. Do **NOT** explain what you changed. Just output the result.
4. The <critique> **OVERRIDES** any conflicting instructions in the <original_prompt>.
   - Example: If original prompt was 'Draft a minimal deployment plan' and critique is 'Include rollback steps', you MUST output the full plan **with** rollback steps included.
5. Your output will be spliced directly into the history as if it were your first and only response.
EOF
fi
trap cleanup EXIT

# Construct the "Frankenstein" Prompt
{
  echo "<refinement_task>"
  echo "    <original_prompt>"
  printf '%s\n' "$ORIG_PROMPT_CONTENT"
  echo "    </original_prompt>"
  echo "    <previous_draft>"
  printf '%s\n' "$ORIG_CONTENT"
  echo "    </previous_draft>"
  echo "    <critique>"
  printf '%s\n' "$INSTRUCTION"
  echo "    </critique>"
  echo "    <instruction>"
  echo "    Rewrite the content of <previous_draft> to satisfy the <critique>."
  echo "    Respond to the <original_prompt> directly, but apply the modifications requested."
  echo "    **Output ONLY the final rewritten content.**"
  echo "    </instruction>"
  echo "</refinement_task>"
} > "$PROMPT_FILE"


echo "Refining response at index $ABS_INDEX (Mode: $ORIG_MODE)..."

if [ "$ORIG_MODE" = "diff" ]; then
    CMD="gen"
else
    CMD="ask"
fi

# 1. Execute inside ephemeral fork
REFINE_LOGIC='
set -e

# 1. Rewind history
aico undo "${REFINE_INDEX}..${REFINE_LAST_INDEX}" >/dev/null

# 2. Generate
cat "$REFINE_PROMPT_FILE" | aico "$REFINE_CMD" $REFINE_MODEL_ARGS --system-prompt "$(cat "$REFINE_SYSPROMPT_FILE")"

# 3. Capture result
aico last --json > "$REFINE_PAYLOAD"
'

# Calculate the absolute end of history for the undo range
LAST_INDEX_DUMP=$(mktemp -t aico_refine_last.XXXXXX)
aico last -1 --json > "$LAST_INDEX_DUMP" 2>/dev/null || true
LAST_INDEX=$(jq -r .pair_index "$LAST_INDEX_DUMP")
rm -f "$LAST_INDEX_DUMP"

# Export all context needed by the inner logic
export REFINE_INDEX="$ABS_INDEX"
export REFINE_LAST_INDEX="$LAST_INDEX"
export REFINE_CMD="$CMD"
export REFINE_MODEL_ARGS="$MODEL_FLAG"
export REFINE_PROMPT_FILE="$PROMPT_FILE"
export REFINE_SYSPROMPT_FILE="$SYSTEM_PROMPT_FILE"
export REFINE_PAYLOAD="$PAYLOAD"

# Execute in ephemeral fork
aico session-fork --ephemeral "$SESSION_NAME" -- sh -c "$REFINE_LOGIC"

# 2. Apply to Main Session
if [ -s "$PAYLOAD" ]; then
    NEW_ASST_ID=$(jq -r .assistant.id "$PAYLOAD")

    if [ "$NEW_ASST_ID" = "null" ] || [ "$ORIG_USER_ID" = "null" ]; then
       echo "Error: IDs could not be resolved. Cannot use surgical history splice."
       exit 1
    fi

    # 1. Undo the old one (soft delete)
    aico undo "$INDEX" >/dev/null

    # 2. Splice the new pair immediately after
    NEXT_INDEX=$((ABS_INDEX + 1))
    aico history-splice "$ORIG_USER_ID" "$NEW_ASST_ID" --at-index "$NEXT_INDEX" >/dev/null

    echo "Refinement complete (Original pair hidden, new pair spliced)."
else
    echo "Refinement failed or returned empty content."
    exit 1
fi
