#!/bin/bash
# schema - inspect database schema
# usage: schema                    (list all tables)
#        schema --schemas          (list all schemas)
#        schema <table>            (inspect specific table, e.g. tinydot.users)
#        schema <schema>.*         (list tables in schema, e.g. tinydot.*)

set -e

# Get connection string
if [ -n "$DATABASE_URL" ]; then
  CONN="$DATABASE_URL"
elif [ -f /workspace/.connection.json ]; then
  CONN=$(cat /workspace/.connection.json | jq -r .connectionString)
else
  echo '{"error": "No database connection. Set DATABASE_URL or provide /workspace/.connection.json"}' >&2
  exit 1
fi

# List schemas
if [ "$1" = "--schemas" ]; then
  psql "$CONN" -t -A -c "
    SELECT json_agg(s)
    FROM (
      SELECT
        schema_name as name,
        (SELECT count(*) FROM information_schema.tables t
         WHERE t.table_schema = s.schema_name AND t.table_type = 'BASE TABLE') as tables
      FROM information_schema.schemata s
      WHERE schema_name NOT IN ('information_schema', 'pg_catalog', 'drizzle')
      ORDER BY schema_name
    ) s
  " | jq '.'
  exit 0
fi

# List tables in specific schema (e.g. tinydot.*)
if [[ "$1" == *.* ]] && [[ "$1" == *\* ]]; then
  SCHEMA=$(echo "$1" | cut -d. -f1)
  psql "$CONN" -t -A -c "
    SELECT json_agg(t)
    FROM (
      SELECT
        table_name as table,
        (SELECT count(*) FROM information_schema.columns c
         WHERE c.table_schema = t.table_schema AND c.table_name = t.table_name) as columns
      FROM information_schema.tables t
      WHERE table_schema = '$SCHEMA'
        AND table_type = 'BASE TABLE'
      ORDER BY table_name
    ) t
  " | jq '.'
  exit 0
fi

if [ -z "$1" ]; then
  # List all tables
  psql "$CONN" -t -A -c "
    SELECT json_agg(t)
    FROM (
      SELECT
        table_schema as schema,
        table_name as table,
        (SELECT count(*) FROM information_schema.columns c
         WHERE c.table_schema = t.table_schema AND c.table_name = t.table_name) as columns
      FROM information_schema.tables t
      WHERE table_schema NOT IN ('information_schema', 'pg_catalog', 'drizzle')
        AND table_type = 'BASE TABLE'
      ORDER BY table_schema, table_name
    ) t
  " | jq '.'
else
  # Inspect specific table
  TABLE="$1"

  # Parse schema.table or just table
  if [[ "$TABLE" == *.* ]]; then
    SCHEMA=$(echo "$TABLE" | cut -d. -f1)
    TNAME=$(echo "$TABLE" | cut -d. -f2)
    SCHEMA_FILTER="AND c.table_schema = '$SCHEMA'"
  else
    TNAME="$TABLE"
    SCHEMA_FILTER=""
  fi

  psql "$CONN" -t -A -c "
    SELECT json_build_object(
      'table', '$TABLE',
      'columns', (
        SELECT json_agg(col ORDER BY col->>'position')
        FROM (
          SELECT json_build_object(
            'name', c.column_name,
            'type', c.data_type,
            'nullable', c.is_nullable = 'YES',
            'default', c.column_default,
            'position', c.ordinal_position
          ) as col
          FROM information_schema.columns c
          WHERE c.table_name = '$TNAME' $SCHEMA_FILTER
        ) cols
      ),
      'primary_key', (
        SELECT json_agg(kcu.column_name)
        FROM information_schema.table_constraints tc
        JOIN information_schema.key_column_usage kcu
          ON tc.constraint_name = kcu.constraint_name
          AND tc.table_schema = kcu.table_schema
        WHERE tc.table_name = '$TNAME' $SCHEMA_FILTER
          AND tc.constraint_type = 'PRIMARY KEY'
      ),
      'foreign_keys', (
        SELECT json_agg(fk)
        FROM (
          SELECT json_build_object(
            'column', kcu.column_name,
            'references', ccu.table_schema || '.' || ccu.table_name || '(' || ccu.column_name || ')'
          ) as fk
          FROM information_schema.table_constraints tc
          JOIN information_schema.key_column_usage kcu
            ON tc.constraint_name = kcu.constraint_name
            AND tc.table_schema = kcu.table_schema
          JOIN information_schema.constraint_column_usage ccu
            ON tc.constraint_name = ccu.constraint_name
            AND tc.table_schema = ccu.constraint_schema
          WHERE tc.table_name = '$TNAME' $SCHEMA_FILTER
            AND tc.constraint_type = 'FOREIGN KEY'
        ) fks
      )
    )
  " | jq '.'
fi
