import { ContinuousRunnerResult, PgClient } from '@bmd-studio/genstack-pg';
import environment from '@bmd-studio/genstack-environment';

const { 
  APP_PREFIX,
  POSTGRES_IDENTITY_ROLE_NAME,
  POSTGRES_PUBLIC_SCHEMA_NAME,
  POSTGRES_HIDDEN_SCHEMA_NAME,
  POSTGRES_DEFAULT_SCHEMA_NAME,
  
  JWT_SECRET,
  JWT_IDENTITY_ID_FIELD,
  JWT_ROLE_FIELD,
  JWT_ALGORITHM,
} = environment.env;

/**
 * Identity checks adapted from: https://www.graphile.org/postgraphile/security/
 * JWT implementation adapted from: https://github.com/michelp/pgjwt and https://github.com/larsar/postgres-jwt-prototype
 */
export default (pgClient: PgClient): ContinuousRunnerResult => {
  return pgClient.query(`

    CREATE OR REPLACE FUNCTION ${POSTGRES_HIDDEN_SCHEMA_NAME}.url_encode(data BYTEA) RETURNS TEXT LANGUAGE SQL AS $$
      SELECT translate(encode(data, 'base64'), E'+/=\n', '-_');
    $$;
    
    CREATE OR REPLACE FUNCTION ${POSTGRES_HIDDEN_SCHEMA_NAME}.url_decode(data TEXT) RETURNS BYTEA LANGUAGE SQL AS $$
      WITH t AS (SELECT translate(data, '-_', '+/')),
          rem AS (SELECT length((SELECT * FROM t)) % 4) -- compute padding size
      SELECT decode(
          (SELECT * FROM t) ||
          CASE WHEN (SELECT * FROM rem) > 0 THEN 
            repeat('=', (4 - (SELECT * FROM rem)))
          ELSE 
            '' 
          END,
        'base64'
      );
    $$;
    
    CREATE OR REPLACE FUNCTION ${POSTGRES_HIDDEN_SCHEMA_NAME}.algorithm_sign(signables TEXT, secret TEXT, algorithm TEXT)
      RETURNS TEXT LANGUAGE SQL AS $$
    WITH
        alg AS (
          SELECT CASE
            WHEN algorithm = 'HS256'
              THEN 'sha256'
            WHEN algorithm = 'HS384'
              THEN 'sha384'
            WHEN algorithm = 'HS512'
              THEN 'sha512'
            ELSE '' END) -- hmac throws error
      SELECT ${POSTGRES_HIDDEN_SCHEMA_NAME}.url_encode(${POSTGRES_DEFAULT_SCHEMA_NAME}.hmac(signables, secret, (SELECT * FROM alg)));
    $$;
    
    CREATE OR REPLACE FUNCTION ${POSTGRES_HIDDEN_SCHEMA_NAME}.sign_jwt(payload JSON, secret TEXT, algorithm TEXT DEFAULT '${JWT_ALGORITHM}')
      RETURNS TEXT LANGUAGE SQL AS $$
      WITH
        header AS (SELECT ${POSTGRES_HIDDEN_SCHEMA_NAME}.url_encode(convert_to('{"alg":"' || algorithm || '","typ":"JWT"}', 'utf8'))),
        payload AS (SELECT ${POSTGRES_HIDDEN_SCHEMA_NAME}.url_encode(convert_to(payload :: TEXT, 'utf8'))),
        signables AS (SELECT (SELECT * FROM header) || '.' || (SELECT * FROM payload))

      SELECT (SELECT * FROM signables) || '.' || ${POSTGRES_HIDDEN_SCHEMA_NAME}.algorithm_sign((SELECT * FROM signables), secret, algorithm);
    $$;
    
    CREATE OR REPLACE FUNCTION ${POSTGRES_HIDDEN_SCHEMA_NAME}.verify_jwt(token TEXT, secret TEXT DEFAULT '${JWT_SECRET}', algorithm TEXT DEFAULT '${JWT_ALGORITHM}')
      RETURNS TABLE(header TEXT, payload TEXT, valid BOOLEAN) LANGUAGE SQL AS $$
      SELECT
        convert_from(${POSTGRES_HIDDEN_SCHEMA_NAME}.url_decode(r [1]), 'utf8') :: TEXT                  AS header,
        convert_from(${POSTGRES_HIDDEN_SCHEMA_NAME}.url_decode(r [2]), 'utf8') :: TEXT                  AS payload,
        r [3] = ${POSTGRES_HIDDEN_SCHEMA_NAME}.algorithm_sign(r [1] || '.' || r [2], secret, algorithm) AS valid
      FROM regexp_split_to_array(token, '\\.') r;
    $$;  
    
    CREATE OR REPLACE FUNCTION ${POSTGRES_PUBLIC_SCHEMA_NAME}.current_identity_id() RETURNS text AS $$
      SELECT NULLIF(current_setting('jwt.claims.${JWT_IDENTITY_ID_FIELD}', true), '')::text;
    $$ LANGUAGE sql STABLE SECURITY DEFINER;

    CREATE OR REPLACE FUNCTION ${POSTGRES_PUBLIC_SCHEMA_NAME}.current_identity_role() RETURNS text AS $$
      SELECT NULLIF(current_setting('jwt.claims.${JWT_ROLE_FIELD}', true), '')::text;
    $$ LANGUAGE sql STABLE SECURITY DEFINER; 

    CREATE OR REPLACE FUNCTION ${POSTGRES_PUBLIC_SCHEMA_NAME}.set_current_identity_by_jwt(token TEXT) RETURNS text AS $$
      DECLARE
          payload JSON;
          identity_id TEXT;
      BEGIN

        -- get the raw JSON payload from the token
        EXECUTE 'SELECT payload::json FROM ${POSTGRES_HIDDEN_SCHEMA_NAME}.verify_jwt($1)' USING token INTO payload;

        -- get the identity ID from the payload
        identity_id := payload->>'identity_id';

        -- TMP debugging
        -- RAISE NOTICE 'Got payload: %', payload;
        -- RAISE NOTICE 'Got identity_id: %', identity_id;

        -- change the role and claim
        EXECUTE 'SELECT set_config(''jwt.claims.identity_id'', ''' || identity_id || ''', false)';
        EXECUTE 'SET role ${APP_PREFIX}_${POSTGRES_IDENTITY_ROLE_NAME}';

        RETURN current_identity_id();
      END      
    $$ LANGUAGE plpgsql VOLATILE;    
  `);
};