#!/bin/bash
# This is free software for the public good of a permacomputer hosted at
# permacomputer.com, an always-on computer by the people, for the people.
# One which is durable, easy to repair, & distributed like tap water
# for machine learning intelligence.
#
# The permacomputer is community-owned infrastructure optimized around
# four values:
#
#   TRUTH      First principles, math & science, open source code freely distributed
#   FREEDOM    Voluntary partnerships, freedom from tyranny & corporate control
#   HARMONY    Minimal waste, self-renewing systems with diverse thriving connections
#   LOVE       Be yourself without hurting others, cooperation through natural law
#
# This software contributes to that vision by making machine learning
# accessible to everyone through a free, open, embeddable chat interface.
# Code is seeds to sprout on any abandoned technology.

# UncloseAI Bash Library - OpenAI-compatible API client with streaming support
# Compatible with vLLM, Ollama, and OpenAI-compatible endpoints
#
# Library Functions:
#   uncloseai_init                      - Initialize client with model discovery
#   uncloseai_list_models               - List discovered models
#   uncloseai_chat <messages_json>      - Non-streaming chat completion
#   uncloseai_chat_stream <messages_json> - Streaming chat completion
#   uncloseai_tts <text> <voice>        - Text-to-speech generation

# Global client state
declare -a UNCLOSEAI_MODEL_IDS
declare -a UNCLOSEAI_MODEL_ENDPOINTS
declare -a UNCLOSEAI_MODEL_MAX_TOKENS
declare -a UNCLOSEAI_TTS_ENDPOINTS

# Initialize client and discover models
uncloseai_init() {
    UNCLOSEAI_MODEL_IDS=()
    UNCLOSEAI_MODEL_ENDPOINTS=()
    UNCLOSEAI_MODEL_MAX_TOKENS=()
    UNCLOSEAI_TTS_ENDPOINTS=()

    # Discover chat/code models from MODEL_ENDPOINT_N
    for i in {1..9999}; do
        local var_name="MODEL_ENDPOINT_$i"
        local endpoint="${!var_name}"

        if [ -z "$endpoint" ]; then
            break
        fi

        local response=$(curl -s "${endpoint}/models")

        # Extract model IDs from JSON, filtering out modelperm entries
        local model_ids=$(echo "$response" | grep -o '"id":"[^"]*"' | sed 's/"id":"//g' | sed 's/"//g' | grep -v "^modelperm-")

        # Add each discovered model to arrays
        while IFS= read -r model_id; do
            if [ -n "$model_id" ]; then
                # Try to extract max_model_len from vLLM response
                local max_tokens=$(echo "$response" | grep -o '"max_model_len":[0-9]*' | head -1 | sed 's/"max_model_len"://g')
                if [ -z "$max_tokens" ]; then
                    max_tokens=8192  # Default for Ollama
                fi

                UNCLOSEAI_MODEL_IDS+=("$model_id")
                UNCLOSEAI_MODEL_ENDPOINTS+=("$endpoint")
                UNCLOSEAI_MODEL_MAX_TOKENS+=("$max_tokens")
            fi
        done <<< "$model_ids"
    done

    # Discover TTS endpoints from TTS_ENDPOINT_N
    for i in {1..9999}; do
        local var_name="TTS_ENDPOINT_$i"
        local endpoint="${!var_name}"

        if [ -z "$endpoint" ]; then
            break
        fi

        UNCLOSEAI_TTS_ENDPOINTS+=("$endpoint")
    done

    return ${#UNCLOSEAI_MODEL_IDS[@]}
}

# List all discovered models
uncloseai_list_models() {
    for i in "${!UNCLOSEAI_MODEL_IDS[@]}"; do
        echo "  - ${UNCLOSEAI_MODEL_IDS[$i]} (max_tokens: ${UNCLOSEAI_MODEL_MAX_TOKENS[$i]})"
    done
}

# Get model index by ID or return 0 for first model
uncloseai_get_model_idx() {
    local model_id="$1"

    if [ -z "$model_id" ]; then
        echo "0"
        return
    fi

    for i in "${!UNCLOSEAI_MODEL_IDS[@]}"; do
        if [ "${UNCLOSEAI_MODEL_IDS[$i]}" == "$model_id" ]; then
            echo "$i"
            return
        fi
    done

    echo "-1"  # Not found
}

# Non-streaming chat completion
# Usage: uncloseai_chat '<messages_json>' [model_id] [max_tokens] [temperature]
uncloseai_chat() {
    local messages_json="$1"
    local model_id="${2:-}"
    local max_tokens="${3:-100}"
    local temperature="${4:-0.7}"

    local model_idx=$(uncloseai_get_model_idx "$model_id")

    if [ "$model_idx" == "-1" ]; then
        echo "ERROR: Model not found"
        return 1
    fi

    local endpoint="${UNCLOSEAI_MODEL_ENDPOINTS[$model_idx]}"
    local model="${UNCLOSEAI_MODEL_IDS[$model_idx]}"

    local response=$(curl -s -X POST "${endpoint}/chat/completions" \
        -H "Content-Type: application/json" \
        -d "{
            \"model\": \"${model}\",
            \"messages\": ${messages_json},
            \"max_tokens\": ${max_tokens},
            \"temperature\": ${temperature},
            \"stream\": false
        }")

    # Extract content using jq if available, otherwise use grep
    if command -v jq &> /dev/null; then
        echo "$response" | jq -r '.choices[0].message.content'
    else
        echo "$response" | grep -o '"content":"[^"]*"' | head -1 | sed 's/"content":"//g' | sed 's/"$//g'
    fi
}

# Streaming chat completion with SSE parsing
# Usage: uncloseai_chat_stream '<messages_json>' [model_id] [max_tokens] [temperature]
uncloseai_chat_stream() {
    local messages_json="$1"
    local model_id="${2:-}"
    local max_tokens="${3:-500}"
    local temperature="${4:-0.7}"

    local model_idx=$(uncloseai_get_model_idx "$model_id")

    if [ "$model_idx" == "-1" ]; then
        echo "ERROR: Model not found"
        return 1
    fi

    local endpoint="${UNCLOSEAI_MODEL_ENDPOINTS[$model_idx]}"
    local model="${UNCLOSEAI_MODEL_IDS[$model_idx]}"

    # Use curl with --no-buffer for line-by-line streaming
    curl -s --no-buffer -X POST "${endpoint}/chat/completions" \
        -H "Content-Type: application/json" \
        -d "{
            \"model\": \"${model}\",
            \"messages\": ${messages_json},
            \"max_tokens\": ${max_tokens},
            \"temperature\": ${temperature},
            \"stream\": true
        }" | while IFS= read -r line; do
            # SSE format: "data: {...}"
            if [[ "$line" =~ ^data:\ (.+)$ ]]; then
                local data="${BASH_REMATCH[1]}"

                # Check for stream termination
                if [ "$data" == "[DONE]" ]; then
                    break
                fi

                # Extract delta content from streaming chunk
                # Use jq if available, otherwise grep
                if command -v jq &> /dev/null; then
                    local content=$(echo "$data" | jq -r '.choices[0].delta.content // empty' 2>/dev/null)
                else
                    local content=$(echo "$data" | grep -o '"content":"[^"\\]*\\*[^"]*"' | sed 's/"content":"//g' | sed 's/"$//g' | sed 's/\\n/\n/g' | sed 's/\\"/"/g')
                fi

                if [ -n "$content" ] && [ "$content" != "null" ]; then
                    printf "%s" "$content"
                fi
            fi
        done
}

# Text-to-speech generation
# Usage: uncloseai_tts <text> [voice] [model] [output_file]
uncloseai_tts() {
    local text="$1"
    local voice="${2:-alloy}"
    local model="${3:-tts-1}"
    local output_file="${4:-/tmp/speech.mp3}"

    if [ ${#UNCLOSEAI_TTS_ENDPOINTS[@]} -eq 0 ]; then
        echo "ERROR: No TTS endpoints available"
        return 1
    fi

    local endpoint="${UNCLOSEAI_TTS_ENDPOINTS[0]}"

    curl -s -X POST "${endpoint}/audio/speech" \
        -H "Content-Type: application/json" \
        -d "{
            \"model\": \"${model}\",
            \"voice\": \"${voice}\",
            \"input\": \"${text}\"
        }" \
        --output "$output_file"

    if [ -f "$output_file" ]; then
        local file_size=$(stat -f%z "$output_file" 2>/dev/null || stat -c%s "$output_file" 2>/dev/null)
        echo "$file_size"
        return 0
    else
        return 1
    fi
}

# Demo usage when run as script
if [ "${BASH_SOURCE[0]}" == "${0}" ]; then
    echo "=== UncloseAI Bash Client (with Streaming) ==="
    echo ""

    # Initialize client with model discovery
    uncloseai_init
    model_count=${#UNCLOSEAI_MODEL_IDS[@]}

    if [ "$model_count" -eq 0 ]; then
        echo "ERROR: No models discovered. Set environment variables:"
        echo "  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc."
        exit 1
    fi

    echo "Discovered $model_count model(s)"
    uncloseai_list_models
    echo ""

    # Non-streaming chat example
    echo "=== Non-Streaming Chat ==="
    messages='[{"role":"system","content":"You are a helpful AI assistant."},{"role":"user","content":"Explain quantum computing in one sentence."}]'

    response=$(uncloseai_chat "$messages")
    echo "Model: ${UNCLOSEAI_MODEL_IDS[0]}"
    echo "Response: $response"
    echo ""

    # Streaming chat example
    echo "=== Streaming Chat ==="

    # Use second model if available, otherwise first
    if [ "$model_count" -ge 2 ]; then
        model_id="${UNCLOSEAI_MODEL_IDS[1]}"
    else
        model_id="${UNCLOSEAI_MODEL_IDS[0]}"
    fi

    echo "Model: $model_id"
    echo "Response: "

    messages='[{"role":"system","content":"You are a coding assistant."},{"role":"user","content":"Write a bash function to check if a port is open"}]'

    uncloseai_chat_stream "$messages" "$model_id" 200
    echo ""
    echo ""

    # TTS example
    if [ ${#UNCLOSEAI_TTS_ENDPOINTS[@]} -gt 0 ]; then
        echo "=== TTS Speech Generation ==="

        text="Hello from UncloseAI Bash client! This demonstrates text to speech with streaming support."
        output_file="/tmp/speech.mp3"

        file_size=$(uncloseai_tts "$text" "alloy" "tts-1" "$output_file")

        if [ $? -eq 0 ]; then
            echo "[OK] Speech file created: $output_file ($file_size bytes)"
            echo ""
        else
            echo "[ERROR] TTS generation failed"
            echo ""
        fi
    fi

    echo "=== Examples Complete ==="
fi