#!/usr/bin/env bash

# File: prole.sh
# Purpose: Manage Ollama environment and run models

# Dotfile to store environment
DOTFILE="$HOME/.prole-ollama"

# Default values
DEFAULT_MODEL="qwen3-coder:30b"
DEFAULT_OLLAMA_HOST="${OLLAMA_HOST:-http://localhost:11434}"

# Function to load existing settings if the dotfile exists
load_dotfile() {
    if [[ -f "$DOTFILE" ]]; then
        # shellcheck source=/dev/null
        source "$DOTFILE"
    else
        OLLAMA_HOST="$DEFAULT_OLLAMA_HOST"
        PROLE_OLLAMA_MODEL="$DEFAULT_MODEL"
    fi
}

# Function to prompt for environment settings and save
set_environment() {
    echo "Current OLLAMA_HOST: ${OLLAMA_HOST}"
    read -rp "Enter OLLAMA_HOST [${OLLAMA_HOST}]: " input_host
    if [[ -n "$input_host" ]]; then
        OLLAMA_HOST="$input_host"
    fi

    echo "Current PROLE_OLLAMA_MODEL: ${PROLE_OLLAMA_MODEL}"
    read -rp "Enter PROLE_OLLAMA_MODEL [${PROLE_OLLAMA_MODEL}]: " input_model
    if [[ -n "$input_model" ]]; then
        PROLE_OLLAMA_MODEL="$input_model"
    fi

    # Save to dotfile
    cat > "$DOTFILE" <<EOL
# Ollama environment for Prole
export OLLAMA_HOST="$OLLAMA_HOST"
export PROLE_OLLAMA_MODEL="$PROLE_OLLAMA_MODEL"
EOL

    # Export to current shell
    export OLLAMA_HOST
    export PROLE_OLLAMA_MODEL

    echo "Environment set and saved to $DOTFILE"
}

# Function to run the model
run_model() {
    set_environment
    echo "Running Ollama model: $PROLE_OLLAMA_MODEL"
    ollama run "$PROLE_OLLAMA_MODEL"
}

# Analyze mode: gather context from ~/dev/prole and pipe into the model
analyze_dev() {
    set_environment
    DEV_ROOT="$HOME/dev/prole"

    if [[ ! -d "$DEV_ROOT" ]]; then
        echo "Development directory $DEV_ROOT does not exist."
        exit 1
    fi

    CONTEXT_FILE=$(mktemp)

    echo "Aggregating development files from $DEV_ROOT..."

    # Collect relevant file types (you can expand this list)
    find "$DEV_ROOT" -type f \
        \( -name "*.py" -o -name "*.js" -o -name "*.ts" -o -name "*.java" -o -name "*.go" -o -name "*.md" -o -name "*.txt" \) \
        -exec cat {} + >> "$CONTEXT_FILE"

    echo "Running analysis with model: $PROLE_OLLAMA_MODEL"

    {
        echo "You are receiving project context from the Prole codebase."
        echo "The following is the combined content of all relevant files:"
        echo
        cat "$CONTEXT_FILE"
        echo
        echo "End of project context. Please analyze this codebase."
    } | ollama run "$PROLE_OLLAMA_MODEL"

    rm -f "$CONTEXT_FILE"
}

# Load previous settings
load_dotfile

# Main command dispatch
case "$1" in
    run)
        run_model
        ;;
    analyze)
        analyze_dev
        ;;
    "" )
        set_environment
        ;;
    *)
        echo "Usage: $0 [run|analyze]"
        exit 1
        ;;
esac