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();

        // PL/pgSQL 动态为所有含 tenant_id 列的表启用 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('ALTER TABLE %I ENABLE ROW LEVEL SECURITY', tbl);
                    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(())
    }

    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('ALTER TABLE %I DISABLE ROW LEVEL SECURITY', tbl);
                END LOOP;
            END;
            $$;
            "#,
        )
        .await?;

        Ok(())
    }
}
