CREATE OR REPLACE FUNCTION audit_trigger()
RETURNS TRIGGER AS $$
DECLARE
    field RECORD;
    user_id INT;
    company_id INT;
    old_value TEXT;
    new_value TEXT;
    entity_id INT;
    column_list TEXT[];
    has_id_column BOOLEAN;
    cache_key TEXT;
BEGIN
    -- Fetch the user ID and company ID from session and convert empty strings to NULL
    user_id := NULLIF(get_current_data('app.lcp_user_id'), '')::INT;
    company_id := NULLIF(get_current_data('app.lcp_company_id'), '')::INT;

    -- Create cache key for this table
    cache_key := 'audit_columns_' || TG_TABLE_SCHEMA || '.' || TG_TABLE_NAME;

    -- Check if we have cached column information
    BEGIN
        column_list := current_setting(cache_key, true)::TEXT[];
    EXCEPTION
        WHEN OTHERS THEN
            column_list := NULL;
    END;

    -- If not cached, fetch and cache column list
    IF column_list IS NULL THEN
        SELECT ARRAY_AGG(column_name ORDER BY ordinal_position)
        INTO column_list
        FROM information_schema.columns
        WHERE table_schema = TG_TABLE_SCHEMA
          AND table_name = TG_TABLE_NAME;

        -- Cache the column list for this session
        BEGIN
            PERFORM set_config(cache_key, column_list::TEXT, false);
        EXCEPTION
            WHEN OTHERS THEN
                -- Ignore cache setting errors
                NULL;
        END;
    END IF;

    -- Check if 'id' column exists in the column list
    has_id_column := 'id' = ANY(column_list);

    -- Extract the `id` value if it exists
    IF has_id_column THEN
        IF TG_OP = 'UPDATE' OR TG_OP = 'DELETE' THEN
            EXECUTE format('SELECT ($1).%I', 'id') INTO entity_id USING OLD;
        ELSE
            EXECUTE format('SELECT ($1).%I', 'id') INTO entity_id USING NEW;
        END IF;
    ELSE
        entity_id := NULL;
    END IF;

    -- Process UPDATE operation
    IF TG_OP = 'UPDATE' THEN
        FOR field IN SELECT unnest(column_list) AS column_name LOOP
            -- Dynamically get the old and new values of the column
            EXECUTE format('SELECT ($1).%I::TEXT', field.column_name) INTO old_value USING OLD;
            EXECUTE format('SELECT ($1).%I::TEXT', field.column_name) INTO new_value USING NEW;

            -- Check if the values are different
            IF old_value IS DISTINCT FROM new_value THEN
                INSERT INTO audit_logs (
                    table_name,
                    entity_id,
                    field_name,
                    old_value,
                    new_value,
                    created_at,
                    user_id,
                    operation_type,
                    company_id
                )
                VALUES (
                    TG_TABLE_NAME,
                    entity_id,
                    field.column_name,
                    old_value,
                    new_value,
                    now(),
                    user_id,
                    'UPDATE',
                    company_id
                );
            END IF;
        END LOOP;

    -- Process DELETE operation
    ELSIF TG_OP = 'DELETE' THEN
        FOR field IN SELECT unnest(column_list) AS column_name LOOP
            -- Dynamically get the old value of the column
            EXECUTE format('SELECT ($1).%I::TEXT', field.column_name) INTO old_value USING OLD;

            -- Insert the old value into the audit_logs for a DELETE operation
            INSERT INTO audit_logs (
                table_name,
                entity_id,
                field_name,
                old_value,
                new_value,
                created_at,
                user_id,
                operation_type,
                company_id
            )
            VALUES (
                TG_TABLE_NAME,
                entity_id,
                field.column_name,
                old_value,
                NULL,
                now(),
                user_id,
                'DELETE',
                company_id
            );
        END LOOP;
    END IF;

    -- Return appropriate record based on operation
    IF TG_OP = 'DELETE' THEN
        RETURN OLD;
    ELSE
        RETURN NEW;
    END IF;

EXCEPTION
    WHEN OTHERS THEN
        RAISE WARNING 'Error in audit_trigger for table %.%: %',
            TG_TABLE_SCHEMA, TG_TABLE_NAME, SQLERRM;
        -- Return appropriate record to allow operation to continue
        IF TG_OP = 'DELETE' THEN
            RETURN OLD;
        ELSE
            RETURN NEW;
        END IF;
END;
$$ LANGUAGE plpgsql;