#!/bin/bash
set -e

POSTGRES_VERSION="${POSTGRES_VERSION:-16}"
MEMORY_MODEL="${MEMORY_MODEL:-all-MiniLM-L6-v2}"

# Install Python if not present
install_python() {
    if command -v python3 &>/dev/null; then
        echo "Python already installed: $(python3 --version)"
        return 0
    fi

    echo "Installing Python..."
    apt-get update
    apt-get install -y python3 python3-pip python3-venv

    echo "Python installed: $(python3 --version)"
}

# Install pgvector extension
install_pgvector() {
    if [ -f "/usr/share/postgresql/${POSTGRES_VERSION}/extension/vector.control" ]; then
        echo "pgvector already installed"
        return 0
    fi

    echo "Installing pgvector..."
    apt-get update
    apt-get install -y postgresql-${POSTGRES_VERSION}-pgvector

    echo "pgvector installed"
}

# Setup Python virtual environment and dependencies
setup_python_env() {
    local skill_dir="$(dirname "$(dirname "$0")")"
    local venv_dir="$skill_dir/.venv"

    if [ -d "$venv_dir" ]; then
        echo "Python venv already exists"
        return 0
    fi

    echo "Creating Python virtual environment..."
    python3 -m venv "$venv_dir"

    echo "Installing Python dependencies..."
    "$venv_dir/bin/pip" install --upgrade pip
    "$venv_dir/bin/pip" install \
        fastapi==0.109.0 \
        uvicorn==0.27.0 \
        sentence-transformers==2.3.1 \
        psycopg2-binary==2.9.9 \
        python-ulid==2.2.0

    echo "Python environment ready"
}

# Pre-download embedding model
download_model() {
    local skill_dir="$(dirname "$(dirname "$0")")"
    local venv_dir="$skill_dir/.venv"

    echo "Pre-downloading embedding model: $MEMORY_MODEL..."
    "$venv_dir/bin/python" -c "
from sentence_transformers import SentenceTransformer
model = SentenceTransformer('$MEMORY_MODEL')
print(f'Model loaded: {model.get_sentence_embedding_dimension()} dimensions')
"
    echo "Model downloaded"
}

# Initialize database schema
init_database() {
    # Wait for postgres to be available
    local retries=30
    while [ $retries -gt 0 ]; do
        if pg_isready -q 2>/dev/null; then
            break
        fi
        echo "Waiting for PostgreSQL..."
        sleep 1
        retries=$((retries - 1))
    done

    if [ $retries -eq 0 ]; then
        echo "PostgreSQL not available - schema will be created on first run"
        return 0
    fi

    # Source postgres env if available
    [ -f /run/vibestack/postgres.env ] && source /run/vibestack/postgres.env

    local pg_user="${POSTGRES_USER:-vibestack}"
    local pg_db="${POSTGRES_DB:-vibestack}"

    echo "Initializing memory schema..."
    psql -U "$pg_user" -d "$pg_db" << 'EOF'
-- Enable pgvector
CREATE EXTENSION IF NOT EXISTS vector;

-- Memories table
CREATE TABLE IF NOT EXISTS memories (
    id TEXT PRIMARY KEY,
    type TEXT NOT NULL,
    content TEXT NOT NULL,
    embedding vector(384),
    metadata JSONB DEFAULT '{}',
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- Indexes
CREATE INDEX IF NOT EXISTS idx_memories_type ON memories(type);
CREATE INDEX IF NOT EXISTS idx_memories_created ON memories(created_at DESC);

-- Vector index (IVFFlat for approximate nearest neighbor)
-- Only create if we have enough rows, otherwise exact search is faster
DO $$
BEGIN
    IF NOT EXISTS (SELECT 1 FROM pg_indexes WHERE indexname = 'idx_memories_embedding') THEN
        CREATE INDEX idx_memories_embedding ON memories
        USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
    END IF;
EXCEPTION WHEN others THEN
    -- Index creation might fail if not enough rows, that's ok
    RAISE NOTICE 'Vector index not created (will use exact search)';
END $$;

SELECT 'Memory schema initialized' as status;
EOF

    echo "Database schema ready"
}

install_python
install_pgvector
setup_python_env
download_model
init_database

echo "Memory skill setup complete"