use sea_orm_migration::prelude::*;

#[derive(DeriveMigrationName)]
pub struct Migration;

#[async_trait::async_trait]
impl MigrationTrait for Migration {
    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
        let conn = manager.get_connection();

        // 替换所有表的 RLS 策略：移除空字符串绕过条件
        // 原策略允许 current_setting(...) = '' 时通过（绕过 RLS），现在要求变量已设置且匹配
        conn.execute_unprepared(
            r#"
            DO $$
            DECLARE
                tbl TEXT;
            BEGIN
                FOR tbl IN
                    SELECT c.table_name FROM information_schema.columns c
                    JOIN information_schema.tables t
                      ON c.table_name = t.table_name AND c.table_schema = t.table_schema
                    WHERE c.column_name = 'tenant_id'
                      AND c.table_schema = 'public'
                      AND t.table_type = 'BASE TABLE'
                    ORDER BY c.table_name
                LOOP
                    EXECUTE format('DROP POLICY IF EXISTS tenant_isolation ON %I', tbl);
                    EXECUTE format(
                        'CREATE POLICY tenant_isolation ON %I USING (
                            current_setting(''app.current_tenant_id'', true) != ''''
                            AND tenant_id = current_setting(''app.current_tenant_id'', true)::uuid
                        )',
                        tbl
                    );
                END LOOP;
            END;
            $$;
            "#,
        )
        .await?;

        Ok(())
    }

    async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
        let conn = manager.get_connection();

        // 回滚：恢复允许空字符串绕过的原策略
        conn.execute_unprepared(
            r#"
            DO $$
            DECLARE
                tbl TEXT;
            BEGIN
                FOR tbl IN
                    SELECT c.table_name FROM information_schema.columns c
                    JOIN information_schema.tables t
                      ON c.table_name = t.table_name AND c.table_schema = t.table_schema
                    WHERE c.column_name = 'tenant_id'
                      AND c.table_schema = 'public'
                      AND t.table_type = 'BASE TABLE'
                    ORDER BY c.table_name
                LOOP
                    EXECUTE format('DROP POLICY IF EXISTS tenant_isolation ON %I', tbl);
                    EXECUTE format(
                        'CREATE POLICY tenant_isolation ON %I USING (
                            current_setting(''app.current_tenant_id'', true) = ''''
                            OR tenant_id = current_setting(''app.current_tenant_id'', true)::uuid
                        )',
                        tbl
                    );
                END LOOP;
            END;
            $$;
            "#,
        )
        .await?;

        Ok(())
    }
}
