"""DuckDB installation script."""

import os
import sys

import duckdb

if len(sys.argv) > 1:
    DUCKDB_HOME_DIR = sys.argv[1]
else:
    DUCKDB_HOME_DIR = os.environ.get('DUCKDB_HOME_DIR', '/var/task/duckdb')

print(f'--- Baking DuckDB Extensions to: {DUCKDB_HOME_DIR} ---')
os.makedirs(DUCKDB_HOME_DIR, exist_ok=True)

# Connect and set home directory
conn = duckdb.connect()
conn.execute(f"SET home_directory='{DUCKDB_HOME_DIR}'")

# Install core extensions
extensions = {'aws', 'encodings', 'excel', 'httpfs'}
for ext in extensions:
    print(f'Installing {ext}...')
    # FORCE_INSTALL ensures version compatibility
    conn.execute(f'FORCE INSTALL {ext}')

# Install community extensions
extensions = {'snowflake'}
for ext in extensions:
    print(f'Installing {ext}...')
    # FORCE_INSTALL ensures version compatibility
    conn.execute(f'FORCE INSTALL {ext} FROM community')

print('--- Extensions installed successfully ---')