#!/usr/bin/env bash
#
# etlworks-agent.sh — Bash client for the Etlworks AI Agent Scripting API
#
# Usage:
#   export ETLWORKS_URL="https://app.etlworks.com"
#   export ETLWORKS_API_KEY="your-api-key"
#
#   # List available tools
#   ./etlworks-agent.sh tools
#
#   # Execute a tool directly
#   ./etlworks-agent.sh tool search_knowledge_base '{"query": "REST connection"}'
#   ./etlworks-agent.sh tool host_list_commands '{}'
#   ./etlworks-agent.sh tool host_cli '{"command": "list connections"}'
#   ./etlworks-agent.sh tool search_templates '{"query": "salesforce"}'
#
#   # Stateless chat (one-shot, complete response)
#   ./etlworks-agent.sh chat "How do I create a REST API connection?"
#
#   # Stateless chat (streaming)
#   ./etlworks-agent.sh chat-stream "How do I schedule a flow?"
#
#   # Stateful session (multi-turn)
#   SESSION_ID=$(./etlworks-agent.sh session-create)
#   ./etlworks-agent.sh session-chat "$SESSION_ID" "How do I create a connection?"
#   ./etlworks-agent.sh session-chat "$SESSION_ID" "Now show me how to schedule it"
#   ./etlworks-agent.sh session-messages "$SESSION_ID"
#
# Requirements: curl, jq (optional, for pretty output)
#

set -euo pipefail

BASE_URL="${ETLWORKS_URL:-http://localhost:8080}"
API_KEY="${ETLWORKS_API_KEY:-}"
API_BASE="${BASE_URL}/rest/v1/ai-agent/api"

if [ -z "$API_KEY" ]; then
    echo "Error: ETLWORKS_API_KEY environment variable is required" >&2
    echo "Generate an API key in Etlworks: Settings > Users > Edit user > API Key" >&2
    exit 1
fi

AUTH_HEADER="Authorization: Bearer ${API_KEY}"

# Check if jq is available for pretty printing
HAS_JQ=false
if command -v jq &>/dev/null; then
    HAS_JQ=true
fi

pretty_json() {
    if $HAS_JQ; then
        jq .
    else
        cat
    fi
}

extract_field() {
    local field="$1"
    if $HAS_JQ; then
        jq -r ".$field"
    else
        # Fallback: basic grep extraction
        grep -o "\"$field\"[[:space:]]*:[[:space:]]*\"[^\"]*\"" | head -1 | sed 's/.*: *"\(.*\)"/\1/'
    fi
}

case "${1:-help}" in

    # ---------------------------------------------------------------
    # Tool operations
    # ---------------------------------------------------------------

    tools)
        echo "--- Available Tools ---"
        curl -s -H "$AUTH_HEADER" "${API_BASE}/tools" | pretty_json
        ;;

    tool)
        TOOL_NAME="${2:?Usage: $0 tool <name> <args_json>}"
        ARGS_JSON="${3:-{}}"
        curl -s -H "$AUTH_HEADER" \
            -H "Content-Type: application/json" \
            -X POST \
            -d "$ARGS_JSON" \
            "${API_BASE}/tools/${TOOL_NAME}/execute" | pretty_json
        ;;

    # ---------------------------------------------------------------
    # Stateless chat
    # ---------------------------------------------------------------

    chat)
        MESSAGE="${2:?Usage: $0 chat <message>}"
        BODY=$(printf '{"message": %s}' "$(echo "$MESSAGE" | jq -Rs . 2>/dev/null || echo "\"$MESSAGE\"")")
        curl -s -H "$AUTH_HEADER" \
            -H "Content-Type: application/json" \
            -X POST \
            -d "$BODY" \
            "${API_BASE}/chat" | pretty_json
        ;;

    chat-stream)
        MESSAGE="${2:?Usage: $0 chat-stream <message>}"
        BODY=$(printf '{"message": %s}' "$(echo "$MESSAGE" | jq -Rs . 2>/dev/null || echo "\"$MESSAGE\"")")
        echo "--- Streaming response ---"
        curl -sN -H "$AUTH_HEADER" \
            -H "Content-Type: application/json" \
            -X POST \
            -d "$BODY" \
            "${API_BASE}/chat/stream" | while IFS= read -r line; do
            # Strip "data: " prefix from SSE events
            if [[ "$line" == data:* ]]; then
                json="${line#data: }"
                type=$(echo "$json" | jq -r '.type' 2>/dev/null || echo "")
                case "$type" in
                    token)
                        # Print tokens inline (no newline) for real-time effect
                        printf '%s' "$(echo "$json" | jq -r '.content' 2>/dev/null)"
                        ;;
                    phase)
                        echo ""
                        echo "[$(echo "$json" | jq -r '.phase' 2>/dev/null)] $(echo "$json" | jq -r '.message' 2>/dev/null)"
                        ;;
                    tool_start)
                        echo ""
                        echo "[tool] $(echo "$json" | jq -r '.tool' 2>/dev/null): $(echo "$json" | jq -r '.message' 2>/dev/null)"
                        ;;
                    tool_end)
                        echo "[tool] $(echo "$json" | jq -r '.tool' 2>/dev/null): done"
                        ;;
                    error)
                        echo ""
                        echo "[ERROR] $(echo "$json" | jq -r '.message' 2>/dev/null)"
                        ;;
                    end)
                        echo ""
                        echo "--- Done ---"
                        ;;
                    usage)
                        echo "[usage] prompt=$(echo "$json" | jq -r '.prompt_tokens' 2>/dev/null) completion=$(echo "$json" | jq -r '.completion_tokens' 2>/dev/null)"
                        ;;
                esac
            fi
        done
        ;;

    # ---------------------------------------------------------------
    # Stateful sessions
    # ---------------------------------------------------------------

    session-create)
        RESULT=$(curl -s -H "$AUTH_HEADER" -X POST "${API_BASE}/sessions")
        SESSION_ID=$(echo "$RESULT" | extract_field "session_id")
        echo "$SESSION_ID"
        ;;

    session-chat)
        SESSION_ID="${2:?Usage: $0 session-chat <session_id> <message>}"
        MESSAGE="${3:?Usage: $0 session-chat <session_id> <message>}"
        BODY=$(printf '{"message": %s}' "$(echo "$MESSAGE" | jq -Rs . 2>/dev/null || echo "\"$MESSAGE\"")")
        curl -s -H "$AUTH_HEADER" \
            -H "Content-Type: application/json" \
            -X POST \
            -d "$BODY" \
            "${API_BASE}/sessions/${SESSION_ID}/chat" | pretty_json
        ;;

    session-chat-stream)
        SESSION_ID="${2:?Usage: $0 session-chat-stream <session_id> <message>}"
        MESSAGE="${3:?Usage: $0 session-chat-stream <session_id> <message>}"
        BODY=$(printf '{"message": %s}' "$(echo "$MESSAGE" | jq -Rs . 2>/dev/null || echo "\"$MESSAGE\"")")
        curl -sN -H "$AUTH_HEADER" \
            -H "Content-Type: application/json" \
            -X POST \
            -d "$BODY" \
            "${API_BASE}/sessions/${SESSION_ID}/chat/stream" | while IFS= read -r line; do
            if [[ "$line" == data:* ]]; then
                json="${line#data: }"
                type=$(echo "$json" | jq -r '.type' 2>/dev/null || echo "")
                case "$type" in
                    token) printf '%s' "$(echo "$json" | jq -r '.content' 2>/dev/null)" ;;
                    end) echo "" ;;
                    error) echo "[ERROR] $(echo "$json" | jq -r '.message' 2>/dev/null)" ;;
                esac
            fi
        done
        ;;

    session-messages)
        SESSION_ID="${2:?Usage: $0 session-messages <session_id>}"
        curl -s -H "$AUTH_HEADER" "${API_BASE}/sessions/${SESSION_ID}/messages" | pretty_json
        ;;

    # ---------------------------------------------------------------
    # Help
    # ---------------------------------------------------------------

    help|--help|-h)
        echo "Etlworks AI Agent Scripting API — Bash Client"
        echo ""
        echo "Environment:"
        echo "  ETLWORKS_URL       Base URL (default: http://localhost:8080)"
        echo "  ETLWORKS_API_KEY   API key (required)"
        echo ""
        echo "Commands:"
        echo "  tools                                   List available tools"
        echo "  tool <name> <args_json>                 Execute a tool"
        echo "  chat <message>                          One-shot chat (complete response)"
        echo "  chat-stream <message>                   One-shot chat (streaming SSE)"
        echo "  session-create                          Create a session (prints session ID)"
        echo "  session-chat <id> <message>             Chat in session (complete response)"
        echo "  session-chat-stream <id> <message>      Chat in session (streaming SSE)"
        echo "  session-messages <id>                   Get session message history"
        echo ""
        echo "Examples:"
        echo "  $0 tool search_knowledge_base '{\"query\": \"REST connection\"}'"
        echo "  $0 chat \"How do I create a connection?\""
        echo "  SESSION=\$($0 session-create)"
        echo "  $0 session-chat \$SESSION \"Create a REST connection to example.com\""
        ;;

    *)
        echo "Unknown command: $1" >&2
        echo "Run '$0 help' for usage" >&2
        exit 1
        ;;
esac