#!/bin/bash

# AfterAgent Hook
# Triggered when Gemini finishes responding to a prompt
# Sends assistant response from hook input and extracts reasoning from transcript

# Source common utilities
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/common.sh"

log "INFO" "AfterAgent hook triggered"

# Read JSON input from stdin
INPUT=$(read_json_input)

log "DEBUG" "Input: $INPUT"

# Extract fields from hook input
SESSION_ID=$(echo "$INPUT" | jq -r '.session_id // empty')
PROMPT_RESPONSE=$(echo "$INPUT" | jq -r '.prompt_response // empty')
TRANSCRIPT_PATH=$(echo "$INPUT" | jq -r '.transcript_path // empty')

EVENTS_SENT=0

# 1. Send assistant response
# Prefer transcript content (accurate) over prompt_response (has streaming duplication bug)
RESPONSE_CONTENT=""

if [ -n "$TRANSCRIPT_PATH" ] && [ -f "$TRANSCRIPT_PATH" ]; then
  # Extract content from the last gemini message in transcript
  RESPONSE_CONTENT=$(jq -r '
    [.messages[] | select(.type == "gemini")] | last |
    if .content | type == "string" then .content
    elif .content | type == "array" then [.content[] | .text // empty] | join("")
    else empty end
  ' "$TRANSCRIPT_PATH" 2>/dev/null)
  log "DEBUG" "Got response from transcript: ${RESPONSE_CONTENT:0:100}..."
fi

# Fall back to prompt_response if transcript extraction failed
if [ -z "$RESPONSE_CONTENT" ] || [ "$RESPONSE_CONTENT" = "null" ]; then
  RESPONSE_CONTENT="$PROMPT_RESPONSE"
  log "DEBUG" "Using prompt_response fallback: ${RESPONSE_CONTENT:0:100}..."
fi

if [ -n "$RESPONSE_CONTENT" ]; then
  EVENT_PAYLOAD=$(jq -n \
    --arg session_id "$SESSION_ID" \
    --arg content "$RESPONSE_CONTENT" \
    --arg hook_event_name "AfterAgent" \
    --arg type "ASSISTANT_RESPONSE" \
    --arg source "DESKTOP" \
    '{
      session_id: $session_id,
      hook_event_name: $hook_event_name,
      type: $type,
      source: $source,
      content: $content
    }')

  log "DEBUG" "Sending assistant response: ${RESPONSE_CONTENT:0:100}..."

  send_to_mcp "event" "$EVENT_PAYLOAD" "$SESSION_ID"

  if [ $? -eq 0 ]; then
    EVENTS_SENT=$((EVENTS_SENT + 1))
    log "INFO" "Assistant response sent successfully"
  else
    log "ERROR" "Failed to send assistant response"
  fi
else
  log "WARN" "No assistant response content available"
fi

# 2. Extract reasoning/thoughts from transcript (secondary - for rich context)
# Gemini transcript format: single JSON with .messages[] array
# Each gemini message may have .thoughts[] with {subject, description, timestamp}
if [ -n "$TRANSCRIPT_PATH" ] && [ -f "$TRANSCRIPT_PATH" ]; then
  log "DEBUG" "Reading transcript for thoughts: $TRANSCRIPT_PATH"

  # Get thoughts from the last gemini message in the transcript
  THOUGHTS=$(jq -r '
    [.messages[] | select(.type == "gemini")] | last |
    .thoughts // [] |
    map(.subject + ": " + .description) |
    join("\n\n")
  ' "$TRANSCRIPT_PATH" 2>/dev/null)

  if [ -n "$THOUGHTS" ] && [ "$THOUGHTS" != "null" ]; then
    REASONING_PAYLOAD=$(jq -n \
      --arg session_id "$SESSION_ID" \
      --arg content "$THOUGHTS" \
      --arg hook_event_name "AfterAgent" \
      --arg type "REASONING" \
      --arg source "DESKTOP" \
      '{
        session_id: $session_id,
        hook_event_name: $hook_event_name,
        type: $type,
        source: $source,
        content: $content
      }')

    log "DEBUG" "Sending reasoning: ${THOUGHTS:0:100}..."

    send_to_mcp "event" "$REASONING_PAYLOAD" "$SESSION_ID"

    if [ $? -eq 0 ]; then
      EVENTS_SENT=$((EVENTS_SENT + 1))
      log "INFO" "Reasoning event sent successfully"
    else
      log "ERROR" "Failed to send reasoning event"
    fi
  else
    log "DEBUG" "No thoughts found in transcript"
  fi
else
  log "DEBUG" "No transcript path available for reasoning extraction"
fi

log "INFO" "AfterAgent hook completed. Events sent: $EVENTS_SENT"
exit 0