#!/bin/bash
set -e

AGENTS_PORT="${AGENTS_PORT:-8800}"
PERSONALITIES_DIR="${PERSONALITIES_DIR:-/personalities}"
WORKSPACES_DIR="${WORKSPACES_DIR:-/workspaces}"
AGENTS_DATA_DIR="${AGENTS_DATA_DIR:-/data/agents}"
CLAUDE_MODEL="${CLAUDE_MODEL:-claude-sonnet-4-20250514}"
INSTANCES_FILE="/var/run/agents/instances.json"

# Initialize instances file
echo '[]' > "$INSTANCES_FILE"

# Port allocation (start from 8900)
NEXT_PORT=8900

# Generate short ID
gen_id() {
    head -c 4 /dev/urandom | xxd -p
}

# Create HTTP handler
create_handler() {
    cat > /tmp/agents_handler.sh << 'HANDLER'
#!/bin/bash

PERSONALITIES_DIR="${PERSONALITIES_DIR:-/personalities}"
WORKSPACES_DIR="${WORKSPACES_DIR:-/workspaces}"
AGENTS_DATA_DIR="${AGENTS_DATA_DIR:-/data/agents}"
INSTANCES_FILE="/var/run/agents/instances.json"
CLAUDE_MODEL="${CLAUDE_MODEL:-claude-sonnet-4-20250514}"

# Read request
read -r request_line
method=$(echo "$request_line" | cut -d' ' -f1)
full_path=$(echo "$request_line" | cut -d' ' -f2)
path=$(echo "$full_path" | cut -d'?' -f1)

content_length=0
while read -r header; do
    header=$(echo "$header" | tr -d '\r')
    [ -z "$header" ] && break
    if [[ "$header" =~ ^[Cc]ontent-[Ll]ength:\ *([0-9]+) ]]; then
        content_length="${BASH_REMATCH[1]}"
    fi
done

body=""
if [ "$content_length" -gt 0 ]; then
    body=$(head -c "$content_length")
fi

send_response() {
    local status="$1"
    local content_type="$2"
    local body="$3"
    local body_length=${#body}

    printf "HTTP/1.1 %s\r\n" "$status"
    printf "Content-Type: %s\r\n" "$content_type"
    printf "Content-Length: %d\r\n" "$body_length"
    printf "Connection: close\r\n"
    printf "\r\n"
    printf "%s" "$body"
}

send_json() {
    send_response "200 OK" "application/json" "$1"
}

send_error() {
    send_response "$1" "application/json" "{\"error\":\"$2\"}"
}

# Generate ID
gen_id() {
    head -c 4 /dev/urandom | xxd -p
}

# Find next available port
find_port() {
    local port=8900
    while [ $port -lt 9000 ]; do
        if ! netstat -tln 2>/dev/null | grep -q ":$port "; then
            echo $port
            return
        fi
        port=$((port + 1))
    done
    echo "0"
}

# List personalities
list_personalities() {
    local result="[]"
    for dir in "$PERSONALITIES_DIR"/*/; do
        [ ! -d "$dir" ] && continue
        local name=$(basename "$dir")
        local desc=""
        local config="$dir/config.yaml"

        if [ -f "$config" ]; then
            desc=$(yq -r '.description // ""' "$config" 2>/dev/null || echo "")
        fi

        result=$(echo "$result" | jq --arg n "$name" --arg d "$desc" \
            '. + [{name: $n, description: $d}]')
    done
    echo "$result"
}

# Get personality details
get_personality() {
    local name="$1"
    local dir="$PERSONALITIES_DIR/$name"

    if [ ! -d "$dir" ]; then
        echo ""
        return
    fi

    local claude_md=""
    local config="{}"

    if [ -f "$dir/CLAUDE.md" ]; then
        claude_md=$(cat "$dir/CLAUDE.md" | jq -Rs '.')
    fi

    if [ -f "$dir/config.yaml" ]; then
        config=$(yq -o=json '.' "$dir/config.yaml" 2>/dev/null || echo "{}")
    fi

    jq -n --arg name "$name" --argjson config "$config" --argjson claude_md "$claude_md" \
        '{name: $name, config: $config, claude_md: $claude_md}'
}

# Spawn instance
spawn_instance() {
    local personality=$(echo "$body" | jq -r '.personality // empty')
    local workspace=$(echo "$body" | jq -r '.workspace // empty')
    local name=$(echo "$body" | jq -r '.name // empty')

    if [ -z "$personality" ]; then
        send_error "400 Bad Request" "personality required"
        return
    fi

    local personality_dir="$PERSONALITIES_DIR/$personality"
    if [ ! -d "$personality_dir" ]; then
        send_error "404 Not Found" "personality not found: $personality"
        return
    fi

    # Generate ID and name
    local id=$(gen_id)
    [ -z "$name" ] && name="${personality}-${id}"

    # Find port
    local port=$(find_port)
    if [ "$port" = "0" ]; then
        send_error "503 Service Unavailable" "no ports available"
        return
    fi

    # Default workspace
    if [ -z "$workspace" ]; then
        if [ -f "$personality_dir/config.yaml" ]; then
            workspace=$(yq -r '.default_workspace // empty' "$personality_dir/config.yaml" 2>/dev/null)
        fi
        [ -z "$workspace" ] && workspace="$WORKSPACES_DIR/$name"
    fi

    mkdir -p "$workspace"

    # Get config
    local model="$CLAUDE_MODEL"
    local max_turns=10
    local tools="Bash,Read,Write,Edit,Glob,Grep"

    if [ -f "$personality_dir/config.yaml" ]; then
        model=$(yq -r ".model // \"$CLAUDE_MODEL\"" "$personality_dir/config.yaml" 2>/dev/null)
        max_turns=$(yq -r '.max_turns // 10' "$personality_dir/config.yaml" 2>/dev/null)
        local config_tools=$(yq -r '.tools // [] | join(",")' "$personality_dir/config.yaml" 2>/dev/null)
        [ -n "$config_tools" ] && tools="$config_tools"
    fi

    # Build system prompt with peer info
    local system_prompt=""
    if [ -f "$personality_dir/CLAUDE.md" ]; then
        system_prompt=$(cat "$personality_dir/CLAUDE.md")
    fi

    # Add peer info
    local instances=$(cat "$INSTANCES_FILE")
    local peers=$(echo "$instances" | jq -r '.[] | "- \(.name): port \(.port), personality: \(.personality)"' 2>/dev/null || echo "")

    if [ -n "$peers" ]; then
        system_prompt="$system_prompt

## Running Peer Agents
$peers

You can communicate with peers via their HTTP API on localhost:{port}."
    fi

    # Start Claude instance
    local log_file="/var/log/agents/${name}.log"
    mkdir -p /var/log/agents

    cd "$workspace"

    # Create wrapper script for this instance
    cat > "/tmp/agent_${id}.sh" << WRAPPER
#!/bin/bash
cd "$workspace"
exec claude --print \
    --model "$model" \
    --max-turns "$max_turns" \
    --system-prompt "\$(cat << 'SYSPROMPT'
$system_prompt
SYSPROMPT
)" \
    --allowedTools "$tools" \
    "\$@"
WRAPPER
    chmod +x "/tmp/agent_${id}.sh"

    # For now, we don't start a long-running process
    # The instance is registered and can be invoked via the agents API

    # Record instance
    local instance=$(jq -n \
        --arg id "$id" \
        --arg name "$name" \
        --arg personality "$personality" \
        --arg workspace "$workspace" \
        --argjson port "$port" \
        --arg status "ready" \
        --arg started_at "$(date -Iseconds)" \
        '{id:$id, name:$name, personality:$personality, workspace:$workspace, port:$port, status:$status, started_at:$started_at}')

    # Update instances file
    local instances=$(cat "$INSTANCES_FILE")
    echo "$instances" | jq --argjson inst "$instance" '. + [$inst]' > "$INSTANCES_FILE"

    send_json "$instance"
}

# List instances
list_instances() {
    cat "$INSTANCES_FILE"
}

# Get instance
get_instance() {
    local id="$1"
    local instance=$(jq -r --arg id "$id" '.[] | select(.id == $id or .name == $id)' "$INSTANCES_FILE")

    if [ -z "$instance" ] || [ "$instance" = "null" ]; then
        send_error "404 Not Found" "instance not found"
        return
    fi

    echo "$instance"
}

# Delete instance
delete_instance() {
    local id="$1"
    local instance=$(jq -r --arg id "$id" '.[] | select(.id == $id or .name == $id)' "$INSTANCES_FILE")

    if [ -z "$instance" ] || [ "$instance" = "null" ]; then
        send_error "404 Not Found" "instance not found"
        return
    fi

    # Kill process if running
    local pid=$(echo "$instance" | jq -r '.pid // empty')
    if [ -n "$pid" ] && [ "$pid" != "null" ]; then
        kill "$pid" 2>/dev/null || true
    fi

    # Remove from instances
    local instances=$(cat "$INSTANCES_FILE")
    echo "$instances" | jq --arg id "$id" '[.[] | select(.id != $id and .name != $id)]' > "$INSTANCES_FILE"

    send_json '{"deleted":true}'
}

# Estimate tokens from text (~4 chars per token)
estimate_tokens() {
    local text="$1"
    local chars=${#text}
    echo $(( (chars + 3) / 4 ))
}

# Store conversation in DB
store_conversation() {
    local conv_id="$1"
    local instance_id="$2"
    local max_context="${3:-200000}"
    local db="$AGENTS_DATA_DIR/agents.db"

    sqlite3 "$db" "INSERT OR IGNORE INTO conversations (id, instance_id, status, max_context) VALUES ('$conv_id', '$instance_id', 'active', $max_context);"
}

# Store message in DB and update token counts
store_message() {
    local conv_id="$1"
    local role="$2"
    local content="$3"
    local db="$AGENTS_DATA_DIR/agents.db"
    local msg_id=$(gen_id)
    local tokens=$(estimate_tokens "$content")

    # Escape single quotes
    local escaped_content=$(echo "$content" | sed "s/'/''/g")

    sqlite3 "$db" "INSERT INTO messages (id, conversation_id, role, content, tokens) VALUES ('$msg_id', '$conv_id', '$role', '$escaped_content', $tokens);"

    # Update conversation token counts
    if [ "$role" = "user" ]; then
        sqlite3 "$db" "UPDATE conversations SET tokens_in = tokens_in + $tokens, total_tokens = total_tokens + $tokens WHERE id = '$conv_id';"
    else
        sqlite3 "$db" "UPDATE conversations SET tokens_out = tokens_out + $tokens, total_tokens = total_tokens + $tokens WHERE id = '$conv_id';"
    fi
}

# Get conversation usage
get_conversation_usage() {
    local conv_id="$1"
    local db="$AGENTS_DATA_DIR/agents.db"

    local row=$(sqlite3 -json "$db" "SELECT id, tokens_in, tokens_out, total_tokens, max_context FROM conversations WHERE id='$conv_id';" 2>/dev/null)

    if [ -z "$row" ] || [ "$row" = "[]" ]; then
        echo "{}"
        return
    fi

    # Calculate usage percentage and cost
    local total=$(echo "$row" | jq -r '.[0].total_tokens // 0')
    local max=$(echo "$row" | jq -r '.[0].max_context // 200000')
    local tokens_in=$(echo "$row" | jq -r '.[0].tokens_in // 0')
    local tokens_out=$(echo "$row" | jq -r '.[0].tokens_out // 0')

    local usage_pct=0
    if [ "$max" -gt 0 ]; then
        usage_pct=$(echo "scale=1; $total * 100 / $max" | bc 2>/dev/null || echo "0")
    fi

    # Estimate cost (Claude Sonnet: $3/1M input, $15/1M output)
    local cost_in=$(echo "scale=4; $tokens_in * 0.000003" | bc 2>/dev/null || echo "0")
    local cost_out=$(echo "scale=4; $tokens_out * 0.000015" | bc 2>/dev/null || echo "0")
    local total_cost=$(echo "scale=4; $cost_in + $cost_out" | bc 2>/dev/null || echo "0")

    local status="ok"
    if [ "${usage_pct%.*}" -ge 90 ]; then
        status="critical"
    elif [ "${usage_pct%.*}" -ge 75 ]; then
        status="warning"
    fi

    jq -n \
        --arg id "$conv_id" \
        --argjson tokens_in "$tokens_in" \
        --argjson tokens_out "$tokens_out" \
        --argjson total "$total" \
        --argjson max "$max" \
        --arg usage_pct "$usage_pct" \
        --arg cost "$total_cost" \
        --arg status "$status" \
        '{conversation_id:$id, tokens_in:$tokens_in, tokens_out:$tokens_out, total_tokens:$total, max_context:$max, usage_percent:($usage_pct|tonumber), estimated_cost_usd:($cost|tonumber), status:$status}'
}

# Get conversation messages
get_conversation() {
    local conv_id="$1"
    local db="$AGENTS_DATA_DIR/agents.db"

    sqlite3 -json "$db" "SELECT * FROM messages WHERE conversation_id='$conv_id' ORDER BY timestamp;"
}

# Invoke instance (run a prompt)
invoke_instance() {
    local id="$1"
    local instance=$(jq -r --arg id "$id" '.[] | select(.id == $id or .name == $id)' "$INSTANCES_FILE")

    if [ -z "$instance" ] || [ "$instance" = "null" ]; then
        send_error "404 Not Found" "instance not found"
        return
    fi

    local prompt=$(echo "$body" | jq -r '.prompt // empty')
    if [ -z "$prompt" ]; then
        send_error "400 Bad Request" "prompt required"
        return
    fi

    local conversation_id=$(echo "$body" | jq -r '.conversation_id // empty')
    local inst_id=$(echo "$instance" | jq -r '.id')
    local workspace=$(echo "$instance" | jq -r '.workspace')

    # Generate new conversation_id if not provided
    local is_new_conversation=false
    if [ -z "$conversation_id" ]; then
        conversation_id=$(gen_id)$(gen_id)
        is_new_conversation=true
    fi

    # Store conversation if new
    if [ "$is_new_conversation" = true ]; then
        store_conversation "$conversation_id" "$inst_id"
    fi

    # Store user message
    store_message "$conversation_id" "user" "$prompt"

    cd "$workspace"

    # Build command - add --resume if continuing conversation
    local result
    if [ "$is_new_conversation" = true ]; then
        result=$("/tmp/agent_${inst_id}.sh" "$prompt" 2>&1) || true
    else
        # Resume existing conversation
        result=$(claude --print \
            --resume "$conversation_id" \
            "$prompt" 2>&1) || true
    fi

    # Store assistant response
    store_message "$conversation_id" "assistant" "$result"

    local result_json=$(echo "$result" | jq -Rs '.')
    send_json "{\"conversation_id\":\"$conversation_id\",\"result\":$result_json}"
}

# List conversations
list_conversations() {
    local db="$AGENTS_DATA_DIR/agents.db"
    sqlite3 -json "$db" "SELECT c.*, i.name as instance_name, i.personality FROM conversations c LEFT JOIN instances i ON c.instance_id = i.id ORDER BY c.started_at DESC LIMIT 50;" 2>/dev/null || echo "[]"
}

# Route requests
case "$method:$path" in
    GET:/health)
        send_json '{"status":"ok"}'
        ;;
    GET:/personalities)
        result=$(list_personalities)
        send_json "$result"
        ;;
    GET:/personalities/*)
        name="${path#/personalities/}"
        result=$(get_personality "$name")
        if [ -z "$result" ]; then
            send_error "404 Not Found" "personality not found"
        else
            send_json "$result"
        fi
        ;;
    POST:/instances)
        spawn_instance
        ;;
    GET:/instances)
        result=$(list_instances)
        send_json "$result"
        ;;
    GET:/instances/*)
        id="${path#/instances/}"
        result=$(get_instance "$id")
        [ -n "$result" ] && send_json "$result"
        ;;
    DELETE:/instances/*)
        id="${path#/instances/}"
        delete_instance "$id"
        ;;
    POST:/instances/*/invoke)
        id=$(echo "$path" | sed 's|/instances/\([^/]*\)/invoke|\1|')
        invoke_instance "$id"
        ;;
    POST:/spawn/*)
        personality="${path#/spawn/}"
        body="{\"personality\":\"$personality\"}"
        spawn_instance
        ;;
    GET:/conversations)
        result=$(list_conversations)
        send_json "$result"
        ;;
    GET:/conversations/*/usage)
        conv_id=$(echo "$path" | sed 's|/conversations/\([^/]*\)/usage|\1|')
        result=$(get_conversation_usage "$conv_id")
        send_json "$result"
        ;;
    GET:/conversations/*)
        conv_id="${path#/conversations/}"
        result=$(get_conversation "$conv_id")
        send_json "$result"
        ;;
    *)
        send_error "404 Not Found" "Endpoints: GET /personalities, GET /instances, POST /instances, GET /conversations, GET /conversations/{id}/usage"
        ;;
esac
HANDLER

    chmod +x /tmp/agents_handler.sh

    # Export vars
    export PERSONALITIES_DIR
    export WORKSPACES_DIR
    export AGENTS_DATA_DIR
    export INSTANCES_FILE
    export CLAUDE_MODEL
}

# Cleanup stale instances on startup
cleanup_stale() {
    echo '[]' > "$INSTANCES_FILE"
}

# Serve
serve() {
    echo "Starting Agents API on port $AGENTS_PORT..."
    echo "Personalities: $PERSONALITIES_DIR"
    echo "Workspaces: $WORKSPACES_DIR"

    exec socat TCP-LISTEN:$AGENTS_PORT,reuseaddr,fork EXEC:"/tmp/agents_handler.sh"
}

cleanup_stale
create_handler
serve