#!/bin/bash
set -e

echo "=== INITIALIZING SOURCE POSTGRESQL ==="

echo "0. Enabling extensions schema, uuid-ossp, pg_trgm, and btree_gin extensions..."
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -c "CREATE SCHEMA IF NOT EXISTS extensions;"
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -c "CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\" SCHEMA extensions;"
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -c "CREATE EXTENSION IF NOT EXISTS \"pg_trgm\";"
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -c "CREATE EXTENSION IF NOT EXISTS \"btree_gin\";"

echo "0b. Creating triage schema and update_updated_at trigger function..."
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -c "CREATE SCHEMA IF NOT EXISTS triage;"
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -c "
CREATE OR REPLACE FUNCTION triage.update_updated_at()
RETURNS TRIGGER AS \$\$
BEGIN
    NEW.updated_at = now();
    RETURN NEW;
END;
\$\$ LANGUAGE plpgsql;"

echo "0c. Creating auth schema, users table, and dummy auth helper functions to satisfy RLS policies..."
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -c "CREATE SCHEMA IF NOT EXISTS auth;"
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -c "
CREATE TABLE IF NOT EXISTS auth.users (
    id uuid PRIMARY KEY
);
CREATE OR REPLACE FUNCTION auth.uid()
RETURNS uuid AS \$\$
BEGIN
    RETURN NULL;
END;
\$\$ LANGUAGE plpgsql;
CREATE OR REPLACE FUNCTION auth.role()
RETURNS text AS \$\$
BEGIN
    RETURN 'authenticated';
END;
\$\$ LANGUAGE plpgsql;"

echo "0d. Creating Supabase dummy roles for RLS policies..."
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -c "
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'authenticated') THEN
        CREATE ROLE authenticated;
    END IF;
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'anon') THEN
        CREATE ROLE anon;
    END IF;
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'service_role') THEN
        CREATE ROLE service_role;
    END IF;
END
\$\$;"

echo "1. Creating hms schema and tables from schema dump..."
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -f /dumps/supabase_schema.sql

echo "1b. Dropping NOT NULL/FOREIGN KEY constraints, converting JSON to TEXT, and relaxing GENERATED ALWAYS columns..."
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -c "
DO \$\$
DECLARE
    r RECORD;
BEGIN
    -- Drop all Foreign Keys
    FOR r IN 
        SELECT tc.constraint_name, tc.table_schema, tc.table_name
        FROM information_schema.table_constraints tc
        WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.table_schema = 'hms'
    LOOP
        BEGIN
            EXECUTE format('ALTER TABLE %I.%I DROP CONSTRAINT %I', r.table_schema, r.table_name, r.constraint_name);
        EXCEPTION WHEN OTHERS THEN
            RAISE NOTICE 'Skipping FK drop for %.%: %', r.table_name, r.constraint_name, SQLERRM;
        END;
    END LOOP;

    -- Drop all NOT NULL constraints (excluding identity columns to prevent errors)
    FOR r IN 
        SELECT table_schema, table_name, column_name
        FROM information_schema.columns
        WHERE table_schema = 'hms' AND is_nullable = 'NO' AND is_identity = 'NO'
    LOOP
        BEGIN
            EXECUTE format('ALTER TABLE %I.%I ALTER COLUMN %I DROP NOT NULL', r.table_schema, r.table_name, r.column_name);
        EXCEPTION WHEN OTHERS THEN
            RAISE NOTICE 'Skipping not-null drop for %.%: %', r.table_name, r.column_name, SQLERRM;
        END;
    END LOOP;

    -- Alter all JSON and JSONB columns to TEXT (to accept python/single-quoted dictionaries during subset import)
    FOR r IN 
        SELECT table_schema, table_name, column_name
        FROM information_schema.columns
        WHERE table_schema = 'hms' AND data_type IN ('json', 'jsonb')
    LOOP
        BEGIN
            EXECUTE format('ALTER TABLE %I.%I ALTER COLUMN %I TYPE text', r.table_schema, r.table_name, r.column_name);
        EXCEPTION WHEN OTHERS THEN
            RAISE NOTICE 'Skipping json-to-text alter for %.%: %', r.table_name, r.column_name, SQLERRM;
        END;
    END LOOP;

    -- Convert all GENERATED ALWAYS columns to standard columns (to allow direct COPY from CSV)
    FOR r IN 
        SELECT table_schema, table_name, column_name, data_type
        FROM information_schema.columns
        WHERE table_schema = 'hms' AND is_generated = 'ALWAYS'
    LOOP
        BEGIN
            EXECUTE format('ALTER TABLE %I.%I DROP COLUMN %I', r.table_schema, r.table_name, r.column_name);
            EXECUTE format('ALTER TABLE %I.%I ADD COLUMN %I %s', r.table_schema, r.table_name, r.column_name, r.data_type);
        EXCEPTION WHEN OTHERS THEN
            RAISE NOTICE 'Skipping generated-column alter for %.%: %', r.table_name, r.column_name, SQLERRM;
        END;
    END LOOP;
END
\$\$;"

echo "2. Importing hms CSV data subsets..."
# We loop through the CSVs and import them using the \copy meta-command
# We use psql's \copy command because it runs client-side inside the container and bypasses filesystem permission issues of COPY
for csv_path in /dumps/hms_subset/*.csv; do
    if [ -f "$csv_path" ]; then
        filename=$(basename "$csv_path")
        tablename="${filename%.csv}"
        echo " - Loading table 'hms.$tablename' from $filename..."
        
        # Truncate table first to prevent duplicate entries if restarted
        psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -c "TRUNCATE hms.\"$tablename\" CASCADE;"
        
        # Run copy command
        psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" -c "\copy hms.\"$tablename\" FROM '$csv_path' DELIMITER ',' CSV HEADER;"
    fi
done

echo "=== SOURCE POSTGRESQL INITIALIZATION COMPLETE ==="