-- ==========================================
-- 04: ROW-LEVEL SECURITY HELPER FUNCTIONS
-- ==========================================
-- This script creates utility functions for RLS management
-- Runs as: postgres (superuser)

\echo '🔒 Creating RLS helper functions...'

-- ==========================================
-- DEVELOPMENT DATABASE
-- ==========================================
\c aurganize_dev

\echo ''
\echo '🔧 Creating helper functions on aurganize_dev...'

-- ==========================================
-- FUNCTION: set_tenant_context
-- Purpose: Sets the current tenant ID in session context
-- ==========================================
CREATE OR REPLACE FUNCTION set_tenant_context(p_tenant_id UUID)
RETURNS VOID AS $$
BEGIN
    -- Set session variable with tenant ID
    -- set_config(setting_name, new_value, is_local)
    --   setting_name: 'app.current_tenant_id' (custom variable)
    --   new_value: p_tenant_id as TEXT
    --   is_local: false = lasts until transaction end
    --             true = lasts until statement end
    PERFORM set_config('app.current_tenant_id', p_tenant_id::TEXT, false);
    
    -- Log for debugging (only visible with client_min_messages = DEBUG)
    RAISE DEBUG 'Tenant context set to: %', p_tenant_id;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Grant execute to application user
GRANT EXECUTE ON FUNCTION set_tenant_context(UUID) TO aurganize_backend_api;

-- Add comment
COMMENT ON FUNCTION set_tenant_context(UUID) IS 
    'Sets the current tenant context for Row-Level Security filtering. Call at the beginning of each request.';

\echo '  ✅ Function set_tenant_context created'

-- ==========================================
-- FUNCTION: get_current_tenant
-- Purpose: Returns the current tenant ID from session context
-- ==========================================
CREATE OR REPLACE FUNCTION get_current_tenant()
RETURNS UUID AS $$
BEGIN
    -- current_setting(setting_name, missing_ok)
    --   setting_name: 'app.current_tenant_id'
    --   missing_ok: true = return NULL if not set (don't error)
    RETURN current_setting('app.current_tenant_id', true)::UUID;
EXCEPTION
    WHEN OTHERS THEN
        -- If conversion fails or any error, return NULL
        RETURN NULL;
END;
$$ LANGUAGE plpgsql STABLE SECURITY DEFINER;

-- Grant execute to application user
GRANT EXECUTE ON FUNCTION get_current_tenant() TO aurganize_backend_api;

-- Add comment
COMMENT ON FUNCTION get_current_tenant() IS 
    'Returns the current tenant ID from session context. Returns NULL if not set.';

\echo '  ✅ Function get_current_tenant created'

-- ==========================================
-- FUNCTION: create_tenant_isolation_policy
-- Purpose: Creates standard RLS policy on a table
-- ==========================================
CREATE OR REPLACE FUNCTION create_tenant_isolation_policy(p_table_name TEXT)
RETURNS VOID AS $$
DECLARE
    v_policy_name TEXT;
BEGIN
    -- Generate policy name
    v_policy_name := p_table_name || '_tenant_isolation';
    
    -- Enable RLS on table
    EXECUTE format('ALTER TABLE %I ENABLE ROW LEVEL SECURITY', p_table_name);
    
    -- Drop policy if exists (for re-running)
    EXECUTE format('DROP POLICY IF EXISTS %I ON %I', v_policy_name, p_table_name);
    
    -- Create policy
    -- USING clause: Controls which rows are visible for SELECT/UPDATE/DELETE
    -- WITH CHECK clause: Controls which rows can be inserted/updated
    EXECUTE format('
        CREATE POLICY %I ON %I
        USING (tenant_id = current_setting(''app.current_tenant_id'', true)::UUID)
        WITH CHECK (tenant_id = current_setting(''app.current_tenant_id'', true)::UUID)
    ', v_policy_name, p_table_name);
    
    RAISE NOTICE '✅ Created RLS policy: % on table: %', v_policy_name, p_table_name;
END;
$$ LANGUAGE plpgsql;

-- Grant execute to application user
GRANT EXECUTE ON FUNCTION create_tenant_isolation_policy(TEXT) TO aurganize_backend_api;

-- Add comment
COMMENT ON FUNCTION create_tenant_isolation_policy(TEXT) IS 
    'Creates a standard tenant isolation RLS policy on the specified table. Table must have a tenant_id column.';

\echo '  ✅ Function create_tenant_isolation_policy created'

-- ==========================================
-- FUNCTION: create_contract_participant_policy
-- Purpose: Creates RLS policy for contract-based access
-- ==========================================
CREATE OR REPLACE FUNCTION create_contract_participant_policy(p_table_name TEXT)
RETURNS VOID AS $$
DECLARE
    v_policy_name TEXT;
BEGIN
    -- Generate policy name
    v_policy_name := p_table_name || '_contract_access';
    
    -- Enable RLS on table
    EXECUTE format('ALTER TABLE %I ENABLE ROW LEVEL SECURITY', p_table_name);
    
    -- Drop policy if exists
    EXECUTE format('DROP POLICY IF EXISTS %I ON %I', v_policy_name, p_table_name);
    
    -- Create policy for contract-related tables
    -- Allows access if tenant is vendor OR consumer in the contract
    EXECUTE format('
        CREATE POLICY %I ON %I
        USING (
            EXISTS (
                SELECT 1 FROM contracts
                WHERE contracts.id = %I.contract_id
                AND (
                    contracts.vendor_tenant_id = current_setting(''app.current_tenant_id'', true)::UUID
                    OR contracts.consumer_tenant_id = current_setting(''app.current_tenant_id'', true)::UUID
                )
            )
        )
        WITH CHECK (
            EXISTS (
                SELECT 1 FROM contracts
                WHERE contracts.id = %I.contract_id
                AND (
                    contracts.vendor_tenant_id = current_setting(''app.current_tenant_id'', true)::UUID
                    OR contracts.consumer_tenant_id = current_setting(''app.current_tenant_id'', true)::UUID
                )
            )
        )
    ', v_policy_name, p_table_name, p_table_name, p_table_name);
    
    RAISE NOTICE '✅ Created contract-based RLS policy: % on table: %', v_policy_name, p_table_name;
END;
$$ LANGUAGE plpgsql;

-- Grant execute to application user
GRANT EXECUTE ON FUNCTION create_contract_participant_policy(TEXT) TO aurganize_backend_api;

-- Add comment
COMMENT ON FUNCTION create_contract_participant_policy(TEXT) IS 
    'Creates an RLS policy for contract-related tables. Grants access if tenant is vendor or consumer. Table must have a contract_id column.';

\echo '  ✅ Function create_contract_participant_policy created'


-- ==========================================
-- FUNCTION: audit_rls_status
-- Purpose: Audit and report RLS status across all tables
-- ==========================================
CREATE OR REPLACE FUNCTION audit_rls_status()
RETURNS TABLE (
    table_name TEXT,
    rls_enabled BOOLEAN,
    policy_count BIGINT,
    status TEXT
) AS $$
BEGIN
    RETURN QUERY
    SELECT 
        t.tablename::TEXT AS table_name,
        t.rowsecurity AS rls_enabled,
        COUNT(p.policyname) AS policy_count,
        CASE 
            WHEN t.rowsecurity AND COUNT(p.policyname) > 0 THEN '✓ Protected'
            WHEN t.rowsecurity AND COUNT(p.policyname) = 0 THEN '⚠ Enabled but no policies'
            ELSE '✗ Not protected'
        END AS status
    FROM pg_tables t
    LEFT JOIN pg_policies p 
        ON p.tablename = t.tablename 
        AND p.schemaname = t.schemaname
    WHERE t.schemaname = 'public'
        AND t.tablename NOT LIKE 'pg_%'  -- Exclude PostgreSQL system tables
        AND t.tablename NOT LIKE 'sql_%'  -- Exclude SQL standard tables
    GROUP BY t.schemaname, t.tablename, t.rowsecurity
    ORDER BY 
        CASE 
            WHEN NOT t.rowsecurity THEN 1  -- Not protected first
            WHEN t.rowsecurity AND COUNT(p.policyname) = 0 THEN 2  -- Enabled but no policies
            ELSE 3  -- Protected
        END,
        t.tablename;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Grant execute to application user and readonly user
GRANT EXECUTE ON FUNCTION audit_rls_status() TO aurganize_backend_api;
GRANT EXECUTE ON FUNCTION audit_rls_status() TO aurganize_readonly;

-- Add comment
COMMENT ON FUNCTION audit_rls_status() IS 
    'Audits and reports Row-Level Security status for all tables in the public schema. Returns table name, RLS enabled status, policy count, and protection status.';

\echo '  ✅ Function audit_rls_status created'


-- ==========================================
-- STAGING DATABASE
-- ==========================================
\c aurganize_staging

\echo ''
\echo '🔧 Creating helper functions on aurganize_staging...'

CREATE OR REPLACE FUNCTION set_tenant_context(p_tenant_id UUID) RETURNS VOID AS $$
BEGIN
    PERFORM set_config('app.current_tenant_id', p_tenant_id::TEXT, false);
    RAISE DEBUG 'Tenant context set to: %', p_tenant_id;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;
GRANT EXECUTE ON FUNCTION set_tenant_context(UUID) TO aurganize_backend_api;

CREATE OR REPLACE FUNCTION get_current_tenant() RETURNS UUID AS $$
BEGIN
    RETURN current_setting('app.current_tenant_id', true)::UUID;
EXCEPTION WHEN OTHERS THEN RETURN NULL;
END;
$$ LANGUAGE plpgsql STABLE SECURITY DEFINER;
GRANT EXECUTE ON FUNCTION get_current_tenant() TO aurganize_backend_api;

CREATE OR REPLACE FUNCTION create_tenant_isolation_policy(p_table_name TEXT) RETURNS VOID AS $$
DECLARE v_policy_name TEXT;
BEGIN
    v_policy_name := p_table_name || '_tenant_isolation';
    EXECUTE format('ALTER TABLE %I ENABLE ROW LEVEL SECURITY', p_table_name);
    EXECUTE format('DROP POLICY IF EXISTS %I ON %I', v_policy_name, p_table_name);
    EXECUTE format('CREATE POLICY %I ON %I USING (tenant_id = current_setting(''app.current_tenant_id'', true)::UUID) WITH CHECK (tenant_id = current_setting(''app.current_tenant_id'', true)::UUID)', v_policy_name, p_table_name);
    RAISE NOTICE '✅ Created RLS policy: % on table: %', v_policy_name, p_table_name;
END;
$$ LANGUAGE plpgsql;
GRANT EXECUTE ON FUNCTION create_tenant_isolation_policy(TEXT) TO aurganize_backend_api;

CREATE OR REPLACE FUNCTION create_contract_participant_policy(p_table_name TEXT) RETURNS VOID AS $$
DECLARE v_policy_name TEXT;
BEGIN
    v_policy_name := p_table_name || '_contract_access';
    EXECUTE format('ALTER TABLE %I ENABLE ROW LEVEL SECURITY', p_table_name);
    EXECUTE format('DROP POLICY IF EXISTS %I ON %I', v_policy_name, p_table_name);
    EXECUTE format('CREATE POLICY %I ON %I USING (EXISTS (SELECT 1 FROM contracts WHERE contracts.id = %I.contract_id AND (contracts.vendor_tenant_id = current_setting(''app.current_tenant_id'', true)::UUID OR contracts.consumer_tenant_id = current_setting(''app.current_tenant_id'', true)::UUID))) WITH CHECK (EXISTS (SELECT 1 FROM contracts WHERE contracts.id = %I.contract_id AND (contracts.vendor_tenant_id = current_setting(''app.current_tenant_id'', true)::UUID OR contracts.consumer_tenant_id = current_setting(''app.current_tenant_id'', true)::UUID)))', v_policy_name, p_table_name, p_table_name, p_table_name);
    RAISE NOTICE '✅ Created contract-based RLS policy: % on table: %', v_policy_name, p_table_name;
END;
$$ LANGUAGE plpgsql;
GRANT EXECUTE ON FUNCTION create_contract_participant_policy(TEXT) TO aurganize_backend_api;

\echo '  ✅ All helper functions created on staging'

-- ==========================================
-- FUNCTION: audit_rls_status
-- Purpose: Audit and report RLS status across all tables
-- ==========================================
CREATE OR REPLACE FUNCTION audit_rls_status()
RETURNS TABLE (table_name TEXT, rls_enabled BOOLEAN, policy_count BIGINT, status TEXT) AS $$
BEGIN
    RETURN QUERY
    SELECT 
        t.tablename::TEXT,
        t.rowsecurity,
        COUNT(p.policyname),
        CASE 
            WHEN t.rowsecurity AND COUNT(p.policyname) > 0 THEN '✓ Protected'
            WHEN t.rowsecurity AND COUNT(p.policyname) = 0 THEN '⚠ Enabled but no policies'
            ELSE '✗ Not protected'
        END
    FROM pg_tables t
    LEFT JOIN pg_policies p ON p.tablename = t.tablename AND p.schemaname = t.schemaname
    WHERE t.schemaname = 'public' AND t.tablename NOT LIKE 'pg_%' AND t.tablename NOT LIKE 'sql_%'
    GROUP BY t.schemaname, t.tablename, t.rowsecurity
    ORDER BY CASE WHEN NOT t.rowsecurity THEN 1 WHEN t.rowsecurity AND COUNT(p.policyname) = 0 THEN 2 ELSE 3 END, t.tablename;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;
GRANT EXECUTE ON FUNCTION audit_rls_status() TO aurganize_backend_api;
GRANT EXECUTE ON FUNCTION audit_rls_status() TO aurganize_readonly;
COMMENT ON FUNCTION audit_rls_status() IS 'Audits and reports Row-Level Security status for all tables.';


-- ==========================================
-- PRODUCTION DATABASE
-- ==========================================
\c aurganize_prod

\echo ''
\echo '🔧 Creating helper functions on aurganize_prod...'

CREATE OR REPLACE FUNCTION set_tenant_context(p_tenant_id UUID) RETURNS VOID AS $$
BEGIN
    PERFORM set_config('app.current_tenant_id', p_tenant_id::TEXT, false);
    RAISE DEBUG 'Tenant context set to: %', p_tenant_id;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;
GRANT EXECUTE ON FUNCTION set_tenant_context(UUID) TO aurganize_backend_api;

CREATE OR REPLACE FUNCTION get_current_tenant() RETURNS UUID AS $$
BEGIN
    RETURN current_setting('app.current_tenant_id', true)::UUID;
EXCEPTION WHEN OTHERS THEN RETURN NULL;
END;
$$ LANGUAGE plpgsql STABLE SECURITY DEFINER;
GRANT EXECUTE ON FUNCTION get_current_tenant() TO aurganize_backend_api;

CREATE OR REPLACE FUNCTION create_tenant_isolation_policy(p_table_name TEXT) RETURNS VOID AS $$
DECLARE v_policy_name TEXT;
BEGIN
    v_policy_name := p_table_name || '_tenant_isolation';
    EXECUTE format('ALTER TABLE %I ENABLE ROW LEVEL SECURITY', p_table_name);
    EXECUTE format('DROP POLICY IF EXISTS %I ON %I', v_policy_name, p_table_name);
    EXECUTE format('CREATE POLICY %I ON %I USING (tenant_id = current_setting(''app.current_tenant_id'', true)::UUID) WITH CHECK (tenant_id = current_setting(''app.current_tenant_id'', true)::UUID)', v_policy_name, p_table_name);
    RAISE NOTICE '✅ Created RLS policy: % on table: %', v_policy_name, p_table_name;
END;
$$ LANGUAGE plpgsql;
GRANT EXECUTE ON FUNCTION create_tenant_isolation_policy(TEXT) TO aurganize_backend_api;

CREATE OR REPLACE FUNCTION create_contract_participant_policy(p_table_name TEXT) RETURNS VOID AS $$
DECLARE v_policy_name TEXT;
BEGIN
    v_policy_name := p_table_name || '_contract_access';
    EXECUTE format('ALTER TABLE %I ENABLE ROW LEVEL SECURITY', p_table_name);
    EXECUTE format('DROP POLICY IF EXISTS %I ON %I', v_policy_name, p_table_name);
    EXECUTE format('CREATE POLICY %I ON %I USING (EXISTS (SELECT 1 FROM contracts WHERE contracts.id = %I.contract_id AND (contracts.vendor_tenant_id = current_setting(''app.current_tenant_id'', true)::UUID OR contracts.consumer_tenant_id = current_setting(''app.current_tenant_id'', true)::UUID))) WITH CHECK (EXISTS (SELECT 1 FROM contracts WHERE contracts.id = %I.contract_id AND (contracts.vendor_tenant_id = current_setting(''app.current_tenant_id'', true)::UUID OR contracts.consumer_tenant_id = current_setting(''app.current_tenant_id'', true)::UUID)))', v_policy_name, p_table_name, p_table_name, p_table_name);
    RAISE NOTICE '✅ Created contract-based RLS policy: % on table: %', v_policy_name, p_table_name;
END;
$$ LANGUAGE plpgsql;
GRANT EXECUTE ON FUNCTION create_contract_participant_policy(TEXT) TO aurganize_backend_api;

\echo '  ✅ All helper functions created on production'

-- ==========================================
-- FUNCTION: audit_rls_status
-- Purpose: Audit and report RLS status across all tables
-- ==========================================
CREATE OR REPLACE FUNCTION audit_rls_status()
RETURNS TABLE (table_name TEXT, rls_enabled BOOLEAN, policy_count BIGINT, status TEXT) AS $$
BEGIN
    RETURN QUERY
    SELECT 
        t.tablename::TEXT,
        t.rowsecurity,
        COUNT(p.policyname),
        CASE 
            WHEN t.rowsecurity AND COUNT(p.policyname) > 0 THEN '✓ Protected'
            WHEN t.rowsecurity AND COUNT(p.policyname) = 0 THEN '⚠ Enabled but no policies'
            ELSE '✗ Not protected'
        END
    FROM pg_tables t
    LEFT JOIN pg_policies p ON p.tablename = t.tablename AND p.schemaname = t.schemaname
    WHERE t.schemaname = 'public' AND t.tablename NOT LIKE 'pg_%' AND t.tablename NOT LIKE 'sql_%'
    GROUP BY t.schemaname, t.tablename, t.rowsecurity
    ORDER BY CASE WHEN NOT t.rowsecurity THEN 1 WHEN t.rowsecurity AND COUNT(p.policyname) = 0 THEN 2 ELSE 3 END, t.tablename;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;
GRANT EXECUTE ON FUNCTION audit_rls_status() TO aurganize_backend_api;
GRANT EXECUTE ON FUNCTION audit_rls_status() TO aurganize_readonly;
COMMENT ON FUNCTION audit_rls_status() IS 'Audits and reports Row-Level Security status for all tables.';


-- ==========================================
-- SUMMARY
-- ==========================================
\echo ''
\echo '=========================================='
\echo '✅ RLS helper functions created!'
\echo '=========================================='
\echo ''
\echo 'Functions created on all databases:'
\echo '  - set_tenant_context(UUID)              Set current tenant'
\echo '  - get_current_tenant()                  Get current tenant'
\echo '  - create_tenant_isolation_policy(TEXT)  Apply RLS to table'
\echo '  - create_contract_participant_policy(TEXT) Contract-based RLS'
\echo ''
\echo 'Usage example:'
\echo '  SELECT set_tenant_context(''550e8400-e29b-41d4-a716-446655440000''::UUID);'
\echo '  SELECT create_tenant_isolation_policy(''users'');'
\echo ''