#!/bin/sh
# SecuBox LocalAI manager - Docker/Podman container support
# Copyright (C) 2025 CyberMind.fr
#
# Uses LocalAI Docker image with all backends included (llama-cpp, etc.)

CONFIG="localai"
CONTAINER_NAME="localai"
LOCALAI_VERSION="v2.25.0"
# Docker image with all backends included
LOCALAI_IMAGE="localai/localai:${LOCALAI_VERSION}-ffmpeg"

# Paths
DATA_PATH="/srv/localai"
MODELS_PATH="/srv/localai/models"

usage() {
	cat <<'EOF'
Usage: localaictl <command>

Commands:
  install            Pull Docker image and setup LocalAI
  check              Run prerequisite checks
  update             Update LocalAI Docker image
  status             Show container and service status
  logs               Show LocalAI logs (use -f to follow)
  shell              Open shell in container

Model Management:
  models             List installed models
  model-install <n>  Install model from preset or URL
  model-remove <n>   Remove installed model

Service Control:
  service-run        Internal: run container under procd
  service-stop       Stop container

API Endpoints (default port 8080):
  /v1/chat/completions   - Chat completion (OpenAI compatible)
  /v1/completions        - Text completion
  /v1/embeddings         - Generate embeddings
  /v1/models             - List available models
  /                      - Web UI

Configuration: /etc/config/localai
EOF
}

require_root() { [ "$(id -u)" -eq 0 ] || { echo "Root required" >&2; exit 1; }; }

log_info() { echo "[INFO] $*"; logger -t localai "$*"; }
log_warn() { echo "[WARN] $*" >&2; logger -t localai -p warning "$*"; }
log_error() { echo "[ERROR] $*" >&2; logger -t localai -p err "$*"; }

uci_get() { uci -q get ${CONFIG}.$1; }
uci_set() { uci set ${CONFIG}.$1="$2" && uci commit ${CONFIG}; }

# Load configuration with defaults
load_config() {
	api_port="$(uci_get main.api_port || echo 8080)"
	api_host="$(uci_get main.api_host || echo 0.0.0.0)"
	data_path="$(uci_get main.data_path || echo /srv/localai)"
	models_path="$(uci_get main.models_path || echo /srv/localai/models)"
	memory_limit="$(uci_get main.memory_limit || echo 2g)"
	threads="$(uci_get main.threads || echo 4)"
	context_size="$(uci_get main.context_size || echo 2048)"
	debug="$(uci_get main.debug || echo 0)"
	cors="$(uci_get main.cors || echo 1)"

	# Docker settings
	docker_image="$(uci_get docker.image || echo $LOCALAI_IMAGE)"

	# Ensure paths exist
	[ -d "$data_path" ] || mkdir -p "$data_path"
	[ -d "$models_path" ] || mkdir -p "$models_path"
}

# =============================================================================
# CONTAINER RUNTIME DETECTION
# =============================================================================

# Detect available container runtime (podman preferred, then docker)
detect_runtime() {
	if command -v podman >/dev/null 2>&1; then
		RUNTIME="podman"
	elif command -v docker >/dev/null 2>&1; then
		RUNTIME="docker"
	else
		RUNTIME=""
	fi
	echo "$RUNTIME"
}

has_runtime() {
	[ -n "$(detect_runtime)" ]
}

run_container() {
	local runtime=$(detect_runtime)
	[ -z "$runtime" ] && { log_error "No container runtime found"; return 1; }
	$runtime "$@"
}

# =============================================================================
# CONTAINER MANAGEMENT
# =============================================================================

container_exists() {
	run_container ps -a --format '{{.Names}}' 2>/dev/null | grep -q "^${CONTAINER_NAME}$"
}

container_running() {
	run_container ps --format '{{.Names}}' 2>/dev/null | grep -q "^${CONTAINER_NAME}$"
}

container_stop() {
	if container_running; then
		log_info "Stopping container..."
		run_container stop "$CONTAINER_NAME" >/dev/null 2>&1 || true
	fi
	if container_exists; then
		run_container rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true
	fi
}

container_pull() {
	load_config
	log_info "Pulling LocalAI Docker image: $docker_image"
	log_info "This may take several minutes (image is ~2-4GB)..."

	if run_container pull "$docker_image"; then
		log_info "Image pulled successfully"
		return 0
	else
		log_error "Failed to pull image"
		return 1
	fi
}

container_run() {
	load_config
	container_stop

	log_info "Starting LocalAI container..."
	log_info "Image: $docker_image"
	log_info "API: http://${api_host}:${api_port}"
	log_info "Models: $models_path"

	# Build environment variables
	local env_args=""
	env_args="$env_args -e LOCALAI_THREADS=$threads"
	env_args="$env_args -e LOCALAI_CONTEXT_SIZE=$context_size"
	[ "$debug" = "1" ] && env_args="$env_args -e LOCALAI_DEBUG=true"
	[ "$cors" = "1" ] && env_args="$env_args -e LOCALAI_CORS=true"

	# Run container in foreground (for procd)
	exec run_container run --rm \
		--name "$CONTAINER_NAME" \
		-p "${api_port}:8080" \
		-v "${models_path}:/models:rw" \
		-v "${data_path}:/build:rw" \
		--memory="$memory_limit" \
		$env_args \
		"$docker_image"
}

container_status() {
	load_config
	local runtime=$(detect_runtime)

	echo "=== LocalAI Status ==="
	echo ""
	echo "Container Runtime: ${runtime:-NOT FOUND}"
	echo ""

	if [ -n "$runtime" ]; then
		if container_running; then
			echo "Container Status: RUNNING"
			echo ""
			run_container ps --filter "name=$CONTAINER_NAME" --format "table {{.Names}}\t{{.Status}}\t{{.Ports}}"
		elif container_exists; then
			echo "Container Status: STOPPED"
		else
			echo "Container Status: NOT CREATED"
		fi
	fi

	echo ""
	echo "=== Configuration ==="
	echo "Image: $docker_image"
	echo "API port: $api_port"
	echo "Data path: $data_path"
	echo "Models path: $models_path"
	echo "Memory limit: $memory_limit"
	echo "Threads: $threads"
	echo "Context size: $context_size"
	echo ""

	# Check API health
	if wget -q -O - "http://127.0.0.1:$api_port/readyz" 2>/dev/null | grep -q "ok"; then
		echo "API Status: HEALTHY"

		# List loaded models via API
		echo ""
		echo "=== Loaded Models (via API) ==="
		local models=$(wget -q -O - "http://127.0.0.1:$api_port/v1/models" 2>/dev/null)
		if [ -n "$models" ]; then
			echo "$models" | jsonfilter -e '@.data[*].id' 2>/dev/null | while read model; do
				echo "  - $model"
			done
		fi
	else
		echo "API Status: NOT RESPONDING"
	fi
}

container_logs() {
	if [ "$1" = "-f" ]; then
		run_container logs -f "$CONTAINER_NAME"
	else
		run_container logs --tail 100 "$CONTAINER_NAME"
	fi
}

container_shell() {
	run_container exec -it "$CONTAINER_NAME" /bin/sh
}

# =============================================================================
# MODEL MANAGEMENT
# =============================================================================

cmd_models() {
	load_config
	echo "=== Installed Models ==="
	echo ""

	if [ -d "$models_path" ]; then
		local count=0
		for model in "$models_path"/*.gguf "$models_path"/*.bin "$models_path"/*.onnx; do
			[ -f "$model" ] || continue
			count=$((count + 1))
			local name=$(basename "$model")
			local size=$(ls -lh "$model" | awk '{print $5}')
			echo "  $count. $name ($size)"
		done

		# Also check for yaml configs (gallery models)
		for yaml in "$models_path"/*.yaml; do
			[ -f "$yaml" ] || continue
			local name=$(basename "$yaml" .yaml)
			echo "  - $name (config)"
		done

		if [ "$count" -eq 0 ]; then
			echo "  No models installed"
			echo ""
			echo "Install a model with:"
			echo "  localaictl model-install tinyllama"
			echo "  localaictl model-install phi2"
		fi
	else
		echo "  Models directory not found: $models_path"
	fi

	echo ""
	echo "=== Available Presets ==="
	echo ""

	# List presets from UCI config
	uci show localai 2>/dev/null | grep "=preset" | while read line; do
		local section=$(echo "$line" | cut -d. -f2 | cut -d= -f1)
		local name=$(uci_get "$section.name")
		local desc=$(uci_get "$section.description")
		local size=$(uci_get "$section.size")
		[ -n "$name" ] && echo "  $name - $desc ($size)"
	done
}

cmd_model_install() {
	load_config
	require_root

	local model_name="$1"
	[ -z "$model_name" ] && { echo "Usage: localaictl model-install <preset-name|url>"; return 1; }

	mkdir -p "$models_path"

	# Check if it's a preset
	local preset_url=""
	local preset_file=""

	# Search presets in UCI
	for section in $(uci show localai 2>/dev/null | grep "=preset" | cut -d. -f2 | cut -d= -f1); do
		local pname=$(uci_get "$section.name")
		if [ "$pname" = "$model_name" ]; then
			preset_url=$(uci_get "$section.url")
			preset_file=$(basename "$preset_url")
			break
		fi
	done

	if [ -n "$preset_url" ]; then
		log_info "Installing preset model: $model_name"
		log_info "URL: $preset_url"
		log_info "This may take a while depending on model size..."

		if wget --show-progress -O "$models_path/$preset_file" "$preset_url"; then
			log_info "Model downloaded: $models_path/$preset_file"

			# Create model config YAML for LocalAI
			cat > "$models_path/$model_name.yaml" << EOF
name: $model_name
backend: llama-cpp
parameters:
  model: $preset_file
context_size: $context_size
threads: $threads
EOF
			log_info "Model config created: $models_path/$model_name.yaml"
			log_info ""
			log_info "Model '$model_name' installed successfully!"
			log_info "Restart LocalAI to load: /etc/init.d/localai restart"
		else
			log_error "Failed to download model"
			return 1
		fi
	elif echo "$model_name" | grep -q "^http"; then
		# Direct URL download
		local filename=$(basename "$model_name")
		log_info "Downloading model from URL..."

		if wget --show-progress -O "$models_path/$filename" "$model_name"; then
			log_info "Model installed: $models_path/$filename"
		else
			log_error "Failed to download model"
			return 1
		fi
	else
		log_error "Unknown model or preset: $model_name"
		echo ""
		echo "Available presets:"
		uci show localai 2>/dev/null | grep "=preset" | while read line; do
			local section=$(echo "$line" | cut -d. -f2 | cut -d= -f1)
			local pname=$(uci_get "$section.name")
			[ -n "$pname" ] && echo "  - $pname"
		done
		return 1
	fi
}

cmd_model_remove() {
	load_config
	require_root

	local model_name="$1"
	[ -z "$model_name" ] && { echo "Usage: localaictl model-remove <model-name>"; return 1; }

	# Find and remove model files
	local found=0
	for ext in gguf bin onnx yaml; do
		local file="$models_path/$model_name.$ext"
		if [ -f "$file" ]; then
			rm -f "$file"
			log_info "Removed: $file"
			found=1
		fi
	done

	# Also try to match partial names (model file might have different name)
	for file in "$models_path"/*"$model_name"*; do
		if [ -f "$file" ]; then
			rm -f "$file"
			log_info "Removed: $file"
			found=1
		fi
	done

	if [ "$found" -eq 0 ]; then
		log_warn "Model not found: $model_name"
	else
		log_info "Restart LocalAI to apply: /etc/init.d/localai restart"
	fi
}

# =============================================================================
# COMMANDS
# =============================================================================

cmd_install() {
	require_root
	load_config

	if ! has_runtime; then
		log_error "No container runtime found!"
		log_error "Install podman or docker first:"
		log_error "  opkg update && opkg install podman"
		exit 1
	fi

	local runtime=$(detect_runtime)
	log_info "Installing LocalAI using $runtime..."
	log_info "Image: $docker_image"

	# Create directories
	mkdir -p "$data_path"
	mkdir -p "$models_path"

	# Pull the image
	container_pull || exit 1

	# Enable service
	uci_set main.enabled '1'
	/etc/init.d/localai enable

	log_info ""
	log_info "LocalAI installed successfully!"
	log_info ""
	log_info "Start with: /etc/init.d/localai start"
	log_info "API endpoint: http://<router-ip>:$api_port/v1"
	log_info "Web UI: http://<router-ip>:$api_port"
	log_info ""
	log_info "Install a model to get started:"
	log_info "  localaictl model-install tinyllama   # Lightweight (669MB)"
	log_info "  localaictl model-install phi2        # Balanced (1.6GB)"
}

cmd_check() {
	load_config

	echo "=== Prerequisite Check ==="
	echo ""

	# Check container runtime
	local runtime=$(detect_runtime)
	if [ -n "$runtime" ]; then
		echo "[OK] Container runtime: $runtime"
		$runtime --version 2>/dev/null | head -1
	else
		echo "[FAIL] No container runtime found"
		echo "       Install: opkg install podman"
	fi
	echo ""

	# Check storage
	local storage_path=$(dirname "$data_path")
	local storage_avail=$(df -h "$storage_path" 2>/dev/null | tail -1 | awk '{print $4}')
	echo "Storage available: $storage_avail (at $storage_path)"
	echo "  Note: LocalAI image requires ~2-4GB"
	echo "        Models require 500MB-8GB each"
	echo ""

	# Check memory
	local mem_total=$(grep MemTotal /proc/meminfo | awk '{print $2}')
	local mem_gb=$((mem_total / 1024 / 1024))
	echo "System memory: ${mem_gb}GB"

	if [ "$mem_gb" -lt 2 ]; then
		echo "[WARN] Low memory! LocalAI requires at least 2GB RAM"
	else
		echo "[OK] Memory sufficient"
	fi
	echo ""

	# Check if image exists
	if [ -n "$runtime" ]; then
		if $runtime images --format '{{.Repository}}:{{.Tag}}' 2>/dev/null | grep -q "localai"; then
			echo "[OK] LocalAI image found"
		else
			echo "[INFO] LocalAI image not downloaded yet"
			echo "       Run: localaictl install"
		fi
	fi
}

cmd_update() {
	require_root
	load_config

	log_info "Updating LocalAI..."

	# Stop if running
	container_stop

	# Pull latest image
	container_pull || exit 1

	# Restart if was enabled
	if [ "$(uci_get main.enabled)" = "1" ]; then
		/etc/init.d/localai restart
	else
		log_info "Update complete. Start manually with: /etc/init.d/localai start"
	fi
}

cmd_status() {
	container_status
}

cmd_logs() {
	container_logs "$@"
}

cmd_shell() {
	if ! container_running; then
		log_error "Container not running. Start with: /etc/init.d/localai start"
		exit 1
	fi
	container_shell
}

cmd_service_run() {
	require_root
	load_config

	if ! has_runtime; then
		log_error "No container runtime found"
		exit 1
	fi

	container_run
}

cmd_service_stop() {
	require_root
	container_stop
}

# Main Entry Point
case "${1:-}" in
	install) shift; cmd_install "$@" ;;
	check) shift; cmd_check "$@" ;;
	update) shift; cmd_update "$@" ;;
	status) shift; cmd_status "$@" ;;
	logs) shift; cmd_logs "$@" ;;
	shell) shift; cmd_shell "$@" ;;
	models) shift; cmd_models "$@" ;;
	model-install) shift; cmd_model_install "$@" ;;
	model-remove) shift; cmd_model_remove "$@" ;;
	service-run) shift; cmd_service_run "$@" ;;
	service-stop) shift; cmd_service_stop "$@" ;;
	help|--help|-h|'') usage ;;
	*) echo "Unknown command: $1" >&2; usage >&2; exit 1 ;;
esac