#!/bin/bash
# AGI Training Pipeline - Main Training Script
# Uses API key to generate training data and improve natural language understanding

set -e

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

# Paths
CHAT_BIN="$AGI_HOME/core/build/bin/agi_chat"
CONV_LOG_DIR="$AGI_HOME/logs/conversations"
TRAINING_DATA="$TRAINING_DATA_DIR/training_corpus.txt"
PROGRESS_FILE="$AGI_HOME/logs/training_progress.json"

mkdir -p "$CONV_LOG_DIR"
mkdir -p "$TRAINING_DATA_DIR"

log "───────────────────────────────────────────────"
log "🧠 AGI TRAINING PIPELINE"
log "───────────────────────────────────────────────"

# Check if agi_chat exists
if [ ! -f "$CHAT_BIN" ]; then
    log "❌ ERROR: agi_chat not found at $CHAT_BIN"
    log "Building agi_chat..."
    cd "$AGI_HOME/core/build"
    cmake .. -DCMAKE_BUILD_TYPE=Release
    make -j$(nproc) agi_chat
    log "✅ agi_chat built successfully"
fi

# Check gossip service
if ! curl -s -m 2 http://localhost:8008/stats > /dev/null; then
    log "⚠️  Gossip service not responding. Training will continue but swarm integration disabled."
    GOSSIP_AVAILABLE=0
else
    SWARM_STATUS=$(curl -s http://localhost:8008/stats)
    log "✅ Gossip service connected"
    log "   $SWARM_STATUS"
    GOSSIP_AVAILABLE=1
fi

# Training conversation prompts (diverse set for natural language understanding)
declare -a TRAINING_PROMPTS=(
    "hello"
    "what is the swarm status?"
    "how many agents are running?"
    "tell me about the migration system"
    "what can you help me with?"
    "show me training capabilities"
    "what is the cluster info?"
    "how does distributed learning work?"
    "explain the fibonacci growth control"
    "what is MAX consensus?"
    "how do agents migrate between nodes?"
    "what is the current goal convergence?"
    "show me system statistics"
    "how can I train you better?"
    "what have you learned so far?"
    "describe your memory system"
    "what is your knowledge base size?"
    "how do you understand natural language?"
    "what commands can you execute?"
    "tell me about the P2P architecture"
)

log ""
log "Starting training with ${#TRAINING_PROMPTS[@]} conversation scenarios..."
log ""

# Function to run single training conversation
run_training_conversation() {
    local session_id=$1
    local prompt=$2
    local conv_log="$CONV_LOG_DIR/session_${session_id}_$(date +%s).log"

    log "Session $session_id: \"$prompt\""

    # Run conversation
    echo -e "$prompt\nexit" | "$CHAT_BIN" > "$conv_log" 2>&1

    # Extract response
    RESPONSE=$(grep "^AGI:" "$conv_log" | head -1 | sed 's/^AGI: //')

    # Log to training corpus
    echo "Q: $prompt" >> "$TRAINING_DATA"
    echo "A: $RESPONSE" >> "$TRAINING_DATA"
    echo "" >> "$TRAINING_DATA"

    log "   Response: ${RESPONSE:0:80}..."
}

# Main training loop
SESSION_COUNT=0
for prompt in "${TRAINING_PROMPTS[@]}"; do
    SESSION_COUNT=$((SESSION_COUNT + 1))
    run_training_conversation $SESSION_COUNT "$prompt"

    # Small delay to prevent overwhelming the system
    sleep 0.5
done

log ""
log "───────────────────────────────────────────────"
log "✅ Training Phase 1 Complete"
log "───────────────────────────────────────────────"
log "Sessions completed: $SESSION_COUNT"
log "Training corpus: $TRAINING_DATA"
log "Conversation logs: $CONV_LOG_DIR"

# Training corpus statistics
CORPUS_LINES=$(wc -l < "$TRAINING_DATA")
CORPUS_WORDS=$(wc -w < "$TRAINING_DATA")
CORPUS_SIZE=$(du -h "$TRAINING_DATA" | cut -f1)

log ""
log "Training Data Statistics:"
log "  Lines: $CORPUS_LINES"
log "  Words: $CORPUS_WORDS"
log "  Size: $CORPUS_SIZE"

# Save progress
cat > "$PROGRESS_FILE" <<EOF
{
  "timestamp": "$(date -u +%Y-%m-%dT%H:%M:%SZ)",
  "sessions_completed": $SESSION_COUNT,
  "corpus_lines": $CORPUS_LINES,
  "corpus_words": $CORPUS_WORDS,
  "credits_used": $CREDITS_USED,
  "max_credits": $MAX_CREDITS,
  "gossip_available": $GOSSIP_AVAILABLE
}
EOF

log ""
log "Progress saved to: $PROGRESS_FILE"

if [ $GOSSIP_AVAILABLE -eq 1 ]; then
    FINAL_STATUS=$(curl -s http://localhost:8008/stats)
    log ""
    log "Final Swarm Status:"
    echo "$FINAL_STATUS" | while read -r line; do
        log "  $line"
    done
fi

log ""
log "───────────────────────────────────────────────"
log "🎓 Training pipeline complete!"
log "───────────────────────────────────────────────"
log ""
log "Next steps:"
log "  1. Review training corpus: $TRAINING_DATA"
log "  2. Analyze conversation logs: $CONV_LOG_DIR"
log "  3. Run interactive session: $CHAT_BIN"
log "  4. Continue training with: $0"
log ""
log "API Credits: $CREDITS_USED / $MAX_CREDITS used"
log ""