#!/bin/bash

# Unified Logging System
# This script provides consistent logging and performance metrics across all thinking mechanisms.

# --- Logging Configuration ---
LOG_DIR=~/tmp/ai_thinking
METRICS_FILE="${LOG_DIR}/performance_metrics.json"
SESSION_LOG="${LOG_DIR}/session_$(date +%Y%m%d_%H%M%S).json"
ERROR_LOG="${LOG_DIR}/errors.log"

# Create logging directory
mkdir -p "${LOG_DIR}"

# --- Error Logging ---
log_error() {
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[ERROR] ${timestamp}: ${message}" >> "${ERROR_LOG}"
    echo "Error: ${message}" >&2
}

log_warning() {
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[WARNING] ${timestamp}: ${message}" >> "${ERROR_LOG}"
    echo "Warning: ${message}" >&2
}

# --- Input Validation Functions ---
validate_file_path() {
    local file_path="$1"
    
    if [ -z "$file_path" ]; then
        return 0  # Empty path is valid (optional)
    fi
    
    if [ ! -f "$file_path" ]; then
        log_error "File not found: $file_path"
        return 1
    fi
    
    if [ ! -r "$file_path" ]; then
        log_error "File not readable: $file_path"
        return 1
    fi
    
    return 0
}

validate_prompt() {
    local prompt="$1"
    local max_length=10000
    
    if [ -z "$prompt" ] || [[ "$prompt" =~ ^[[:space:]]*$ ]]; then
        log_error "Empty or whitespace-only prompt"
        return 1
    fi
    
    if [ ${#prompt} -gt $max_length ]; then
        log_error "Prompt too long (${#prompt} chars, max: $max_length)"
        return 1
    fi
    
    # Basic sanitization - remove potentially dangerous characters
    local sanitized=$(echo "$prompt" | sed 's/[<>"'\''&]/g' 2>/dev/null || echo "$prompt")
    if [ "$sanitized" != "$prompt" ]; then
        log_warning "Prompt contained special characters that were sanitized"
        echo "$sanitized"
        return 0
    fi
    
    echo "$prompt"
    return 0
}

# --- Model Validation ---
validate_model() {
    local model="$1"
    local fallback_model="$2"
    
    if ! command -v ollama >/dev/null 2>&1; then
        log_error "Ollama not found in PATH"
        return 1
    fi
    
    if ! ollama list | grep -q "$model"; then
        log_warning "Model '$model' not available"
        if [ -n "$fallback_model" ] && ollama list | grep -q "$fallback_model"; then
            log_warning "Falling back to '$fallback_model'"
            echo "$fallback_model"
            return 0
        else
            log_error "No fallback model available"
            return 1
        fi
    fi
    
    echo "$model"
    return 0
}

# --- Timing Functions ---
start_timer() {
    local session_id="$1"
    local mechanism="$2"
    local start_time=$(date +%s.%N)
    
    # Store start time
    echo "$start_time" > "/tmp/${session_id}_start"
    
    # Log session start
    log_session_start "$session_id" "$mechanism" "$start_time"
}

end_timer() {
    local session_id="$1"
    local mechanism="$2"
    local end_time=$(date +%s.%N)
    local start_time=$(cat "/tmp/${session_id}_start" 2>/dev/null || echo "$end_time")
    
    # Calculate duration
    local duration=$(echo "$end_time - $start_time" | bc -l 2>/dev/null || echo "0")
    
    # Log session end
    log_session_end "$session_id" "$mechanism" "$end_time" "$duration"
    
    # Clean up
    rm -f "/tmp/${session_id}_start"
    
    echo "$duration"
}

# --- Session Logging ---
log_session_start() {
    local session_id="$1"
    local mechanism="$2"
    local start_time="$3"
    
    cat > "${SESSION_LOG}" << EOF
{
  "session_id": "${session_id}",
  "mechanism": "${mechanism}",
  "start_time": "${start_time}",
  "prompt": "${PROMPT:-""}",
  "status": "started"
}
EOF
}

log_session_end() {
    local session_id="$1"
    local mechanism="$2"
    local end_time="$3"
    local duration="$4"
    
    # Update session log
    cat > "${SESSION_LOG}" << EOF
{
  "session_id": "${session_id}",
  "mechanism": "${mechanism}",
  "start_time": "$(cat "${SESSION_LOG}" | jq -r '.start_time' 2>/dev/null || echo "")",
  "end_time": "${end_time}",
  "duration": "${duration}",
  "prompt": "${PROMPT:-""}",
  "status": "completed"
}
EOF
    
    # Update metrics file
    update_metrics "$mechanism" "$duration"
}

# --- Metrics Management ---
update_metrics() {
    local mechanism="$1"
    local duration="$2"
    
    # Create metrics file if it doesn't exist
    if [ ! -f "${METRICS_FILE}" ]; then
        cat > "${METRICS_FILE}" << EOF
{
  "total_sessions": 0,
  "mechanisms": {},
  "average_durations": {}
}
EOF
    fi
    
    # Update metrics using jq (if available) or simple text processing
    if command -v jq >/dev/null 2>&1; then
        # Use jq for proper JSON handling
        local temp_file=$(mktemp)
        jq --arg mechanism "$mechanism" --arg duration "$duration" '
            .total_sessions += 1 |
            .mechanisms[$mechanism] = (.mechanisms[$mechanism] // 0) + 1 |
            .average_durations[$mechanism] = (
                (.average_durations[$mechanism] // 0) * (.mechanisms[$mechanism] - 1) + ($duration | tonumber)
            ) / .mechanisms[$mechanism]
        ' "${METRICS_FILE}" > "$temp_file"
        mv "$temp_file" "${METRICS_FILE}"
    else
        # Fallback: simple text-based metrics
        echo "$(date +%Y%m%d_%H%M%S),${mechanism},${duration}" >> "${LOG_DIR}/simple_metrics.csv"
    fi
}

# --- Utility Functions ---
generate_session_id() {
    echo "session_$(date +%Y%m%d_%H%M%S)_$$"
}

get_metrics_summary() {
    if [ -f "${METRICS_FILE}" ]; then
        echo "=== Performance Metrics ==="
        if command -v jq >/dev/null 2>&1; then
            jq -r '.mechanisms | to_entries[] | "\(.key): \(.value) sessions"' "${METRICS_FILE}"
            echo ""
            jq -r '.average_durations | to_entries[] | "\(.key): \(.value | tonumber | floor)s average"' "${METRICS_FILE}"
        else
            echo "Metrics available in: ${METRICS_FILE}"
        fi
    else
        echo "No metrics available yet."
    fi
}

# --- Export Functions for Other Scripts ---
export -f start_timer
export -f end_timer
export -f log_session_start
export -f log_session_end
export -f update_metrics
export -f generate_session_id
export -f get_metrics_summary
export -f log_error
export -f log_warning
export -f validate_file_path
export -f validate_prompt
export -f validate_model

# Alternative export method for compatibility
if [ -n "$BASH_VERSION" ]; then
    export start_timer end_timer log_session_start log_session_end update_metrics
    export generate_session_id get_metrics_summary log_error log_warning
    export validate_file_path validate_prompt validate_model
fi