-- ============================================================================
-- MagicAuth Postgres Installation Script
-- ============================================================================
--
-- This script sets up a tenant's Postgres database to validate JWTs issued
-- by the MagicAuth service using the tenant's public key (RS256).
--
-- Prerequisites:
--   - plpython3u extension (for JWT verification)
--   - PyJWT Python package installed on the PostgreSQL server:
--       pip install pyjwt cryptography
--
-- Usage:
--   psql -d your_database -f install.sql
--
-- After running this script:
--   1. Insert your tenant's public key into auth.trusted_tenants
--   2. Call auth.set_jwt() at the start of each session with the JWT
--   3. Use auth.email(), auth.role(), auth.tenant_id() in RLS
--
-- ============================================================================

-- Create auth schema
CREATE SCHEMA IF NOT EXISTS auth;

-- Enable required extensions
CREATE EXTENSION IF NOT EXISTS plpython3u;  -- for RS256 JWT verification

-- ============================================================================
-- TABLES
-- ============================================================================

-- Trusted tenants table: stores public keys for JWT verification
-- Each tenant registers their public key here (obtained from /v1/register-tenant)
CREATE TABLE IF NOT EXISTS auth.trusted_tenants (
    tenant_id TEXT PRIMARY KEY,
    public_key_pem TEXT NOT NULL,
    created_at TIMESTAMPTZ DEFAULT now()
);

-- App users table: stores user information and permissions
-- Tenants manage this table to define which emails have access
CREATE TABLE IF NOT EXISTS app_users (
    email TEXT PRIMARY KEY,
    role TEXT NOT NULL DEFAULT 'user',
    is_active BOOLEAN DEFAULT true
);

-- ============================================================================
-- JWT HELPER FUNCTIONS
-- ============================================================================

-- Decode JWT payload without verification
-- Used only as input to verify_jwt to extract tenant_id for public key lookup
-- WARNING: Do not trust this output for authorization - always use verify_jwt
CREATE OR REPLACE FUNCTION auth.decode_jwt_payload_unsafe(token TEXT)
RETURNS JSONB
LANGUAGE sql
IMMUTABLE
STRICT
AS $$
    SELECT convert_from(
        decode(
            rpad(
                translate((string_to_array(token, '.'))[2], '-_', '+/'),
                4 * ceil(length((string_to_array(token, '.'))[2]) / 4.0)::int,
                '='
            ),
            'base64'
        ),
        'UTF8'
    )::jsonb;
$$;

COMMENT ON FUNCTION auth.decode_jwt_payload_unsafe(TEXT) IS
    'Decodes JWT payload WITHOUT signature verification. Never use for authorization.';

-- ============================================================================
-- JWT VERIFICATION (RS256 via PL/Python)
-- ============================================================================

-- Verify JWT signature and expiration using the tenants public key
-- Returns the verified claims if valid, raises exception if invalid
-- Uses PyJWT library for RS256 verification
CREATE OR REPLACE FUNCTION auth.verify_jwt(token TEXT)
RETURNS JSONB
LANGUAGE plpython3u
SECURITY DEFINER
AS $$
import json

# Validate input
if not token:
    plpy.error('JWT token is required')

# First, decode without verification to get tenant_id
try:
    import jwt
    unverified = jwt.decode(token, options={"verify_signature": False})
except Exception as e:
    plpy.error(f'Malformed JWT: {str(e)}')

tenant_id = unverified.get('tenant_id')
if not tenant_id:
    plpy.error('JWT missing tenant_id claim')

# Look up tenant's public key from trusted_tenants (use prepared statement for safety)
plan = plpy.prepare("SELECT public_key_pem FROM auth.trusted_tenants WHERE tenant_id = $1", ["text"])
result = plpy.execute(plan, [tenant_id])

if len(result) == 0:
    plpy.error(f'Unknown tenant: {tenant_id}')

public_key = result[0]['public_key_pem']

# Verify the JWT with the public key
try:
    # PyJWT automatically verifies signature and expiration
    verified = jwt.decode(
        token,
        public_key,
        algorithms=['RS256'],
        options={
            "require": ["exp", "iat", "tenant_id", "email"],
            "verify_exp": True,
            "verify_iat": True
        }
    )
    return json.dumps(verified)
except jwt.ExpiredSignatureError:
    plpy.error('JWT has expired')
except jwt.InvalidSignatureError:
    plpy.error('Invalid JWT signature')
except jwt.InvalidTokenError as e:
    plpy.error(f'Invalid JWT: {str(e)}')
$$;

COMMENT ON FUNCTION auth.verify_jwt(TEXT) IS
    'Verifies JWT signature using tenant public key and checks expiration. Returns claims if valid.';

-- ============================================================================
-- SESSION CLAIMS MANAGEMENT
-- ============================================================================

-- Set session claims from a verified JWT
-- This is the main entry point - call this at the start of each request
CREATE OR REPLACE FUNCTION auth.set_jwt(token TEXT)
RETURNS VOID
LANGUAGE plpgsql
SECURITY DEFINER
AS $$
DECLARE
    claims JSONB;
    u RECORD;
BEGIN
    -- Verify JWT and get claims
    claims := auth.verify_jwt(token);

    -- Set base claims from JWT
    PERFORM set_config('request.jwt.claim.email', claims->>'email', true);
    PERFORM set_config('request.jwt.claim.tenant_id', claims->>'tenant_id', true);
    PERFORM set_config('request.jwt.claim.sub', claims->>'sub', true);
    PERFORM set_config('request.jwt.claim.iss', claims->>'iss', true);
    PERFORM set_config('request.jwt.claim.iat', claims->>'iat', true);
    PERFORM set_config('request.jwt.claim.exp', claims->>'exp', true);

    -- Look up user in app_users to get role
    SELECT * INTO u
    FROM app_users
    WHERE email_address = claims->>'email'
      AND is_active = true;

    IF FOUND THEN
        PERFORM set_config('request.jwt.claim.role', u.role, true);
        PERFORM set_config('request.jwt.claim.is_registered', 'true', true);
    ELSE
        -- User authenticated via email but not in app_users
        -- Set flag so application can handle registration
        PERFORM set_config('request.jwt.claim.role', '', true);
        PERFORM set_config('request.jwt.claim.is_registered', 'false', true);
    END IF;
END;
$$;

COMMENT ON FUNCTION auth.set_jwt(TEXT) IS
    'Validates JWT and sets session claims for use in RLS policies. Call at start of each request.';

-- Clear session claims (call at end of request or on error)
CREATE OR REPLACE FUNCTION auth.clear_jwt()
RETURNS VOID
LANGUAGE plpgsql
SECURITY DEFINER
AS $$
BEGIN
    PERFORM set_config('request.jwt.claim.email', '', true);
    PERFORM set_config('request.jwt.claim.tenant_id', '', true);
    PERFORM set_config('request.jwt.claim.role', '', true);
    PERFORM set_config('request.jwt.claim.sub', '', true);
    PERFORM set_config('request.jwt.claim.iss', '', true);
    PERFORM set_config('request.jwt.claim.iat', '', true);
    PERFORM set_config('request.jwt.claim.exp', '', true);
    PERFORM set_config('request.jwt.claim.is_registered', '', true);
END;
$$;

COMMENT ON FUNCTION auth.clear_jwt() IS
    'Clears all JWT session claims. Call at end of request for connection pooling.';

-- ============================================================================
-- CLAIM GETTER FUNCTIONS
-- ============================================================================
-- Use these in RLS policies and application queries

CREATE OR REPLACE FUNCTION auth.email()
RETURNS TEXT
LANGUAGE sql
STABLE
AS $$
    SELECT nullif(current_setting('request.jwt.claim.email', true), '');
$$;

CREATE OR REPLACE FUNCTION auth.role()
RETURNS TEXT
LANGUAGE sql
STABLE
AS $$
    SELECT nullif(current_setting('request.jwt.claim.role', true), '');
$$;

CREATE OR REPLACE FUNCTION auth.tenant_id()
RETURNS TEXT
LANGUAGE sql
STABLE
AS $$
    SELECT nullif(current_setting('request.jwt.claim.tenant_id', true), '');
$$;

CREATE OR REPLACE FUNCTION auth.is_registered()
RETURNS BOOLEAN
LANGUAGE sql
STABLE
AS $$
    SELECT current_setting('request.jwt.claim.is_registered', true) = 'true';
$$;

-- ============================================================================
-- CONVENIENCE FUNCTIONS
-- ============================================================================

-- Check if user has a specific role
CREATE OR REPLACE FUNCTION auth.has_role(required_role TEXT)
RETURNS BOOLEAN
LANGUAGE sql
STABLE
AS $$
    SELECT auth.role() = required_role;
$$;

-- Check if user has any of the specified roles
CREATE OR REPLACE FUNCTION auth.has_any_role(VARIADIC roles TEXT[])
RETURNS BOOLEAN
LANGUAGE sql
STABLE
AS $$
    SELECT auth.role() = ANY(roles);
$$;

-- Get the authenticated users record from app_users
CREATE OR REPLACE FUNCTION auth.current_user_record()
RETURNS app_users
LANGUAGE sql
STABLE
SECURITY DEFINER
AS $$
    SELECT * FROM app_users
    WHERE email_address = auth.email()
    LIMIT 1;
$$;

-- ============================================================================
-- COMMENTS
-- ============================================================================

COMMENT ON SCHEMA auth IS 'MagicAuth JWT validation and session management';
COMMENT ON TABLE auth.trusted_tenants IS 'Stores trusted tenant public keys for RS256 JWT verification';
COMMENT ON TABLE app_users IS 'Application users with roles and permissions';

COMMENT ON FUNCTION auth.email() IS 'Returns authenticated user email from JWT';
COMMENT ON FUNCTION auth.role() IS 'Returns authenticated user role (NULL if not registered)';
COMMENT ON FUNCTION auth.tenant_id() IS 'Returns tenant_id from JWT';
COMMENT ON FUNCTION auth.is_registered() IS 'Returns true if authenticated user exists in app_users';

-- ============================================================================
-- EXAMPLE: REGISTERING A TRUSTED TENANT
-- ============================================================================
--
-- After obtaining the public key from POST /v1/register-tenant:
--
--   INSERT INTO auth.trusted_tenants (tenant_id, public_key_pem)
--   VALUES (
--       '550e8400-e29b-41d4-a716-446655440000',
--       '-----BEGIN PUBLIC KEY-----
--   MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA...
--   -----END PUBLIC KEY-----'
--   );
--
-- ============================================================================
-- EXAMPLE: USING IN RLS POLICIES
-- ============================================================================
--
--   -- Users can only see their own data
--   CREATE POLICY user_isolation ON my_table
--       FOR ALL
--       USING (email_address = auth.email());
--
-- ============================================================================
-- EXAMPLE: CLIENT USAGE
-- ============================================================================
--
--   BEGIN;
--   SELECT auth.set_jwt('eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9...');
--
--   -- Now RLS policies can use auth.* functions
--   SELECT * FROM my_table;
--
--   -- Optionally check if user is registered
--   SELECT auth.is_registered();  -- returns false for new users
--
--   -- Register new user if needed
--   INSERT INTO app_users (email, role)
--   VALUES (auth.email(), 'user')
--   ON CONFLICT (email) DO NOTHING;
--
--   COMMIT;
--
-- ============================================================================