#!/bin/bash
# Install pi-rtk (Token Reduction Kit) extension to pi agent
# Based on https://github.com/mcowger/pi-rtk

set -e

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

PI_DIR="${PI_DIR:-$HOME/.pi/agent}"
EXT_DIR="$PI_DIR/extensions"
SETTINGS_FILE="$PI_DIR/settings.json"

RTK_SOURCE="https://github.com/mcowger/pi-rtk.git"
RTK_DIR="$EXT_DIR/pi-rtk"

echo -e "${BLUE}╔══════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║       pi-rtk Installation                          ║${NC}"
echo -e "${BLUE}╚══════════════════════════════════════════════════════╝${NC}"
echo ""

mkdir -p "$EXT_DIR"

# Clone or update pi-rtk
if [ -d "$RTK_DIR/.git" ]; then
    echo -e "${YELLOW}Updating existing pi-rtk...${NC}"
    cd "$RTK_DIR" && git pull --ff origin main
else
    echo -e "${YELLOW}Cloning pi-rtk from GitHub...${NC}"
    git clone --depth 1 "$RTK_SOURCE" "$RTK_DIR"
fi

# Update settings.json
echo -e "${YELLOW}Updating settings...${NC}"
python3 << 'PYEOF'
import json
import sys

settings_file = sys.argv[1]
pkg_name = sys.argv[2]

try:
    with open(settings_file, 'r') as f:
        settings = json.load(f)
except:
    settings = {}

if 'packages' not in settings:
    settings['packages'] = []

# Remove old pipeline extensions, keep only rtk
settings['packages'] = [p for p in settings['packages'] if 'pi-rtk' in p]
if pkg_name not in settings['packages']:
    settings['packages'].append(pkg_name)

with open(settings_file, 'w') as f:
    json.dump(settings, f, indent=2)

print(f"Added {pkg_name} to packages")
PYEOF
    "$SETTINGS_FILE" "extensions/pi-rtk/index.ts"

echo ""
echo -e "${GREEN}╔══════════════════════════════════════════════════════╗${NC}"
echo -e "${GREEN}║       Installation Complete!                       ║${NC}"
echo -e "${GREEN}╚══════════════════════════════════════════════════════╝${NC}"
echo ""
echo -e "RTK commands:"
echo -e "  ${BLUE}/rtk-on${NC} / ${BLUE}/rtk-off${NC}    - Enable/disable RTK"
echo -e "  ${BLUE}/rtk-stats${NC}            - Show token savings"
echo -e "  ${BLUE}/rtk-what${NC}             - Show current config"
echo -e "  ${BLUE}/rtk-toggle-*${NC}          - Toggle features (ansiStripping, truncation, etc.)"
echo ""
echo -e "${GREEN}Restart pi agent to load the extension${NC}"