"""Add tenant_id/node_id columns, auto-populate trigger, and RLS to hierarchical_nodes.

The data_hierarchical_nodes table is created by LlamaIndex PostgresDocumentStore.
LlamaIndex only manages (id, key, namespace, value) — it doesn't know about tenant_id/node_id.

Strategy:
1. Add dedicated tenant_id and node_id UUID columns (NOT NULL with default nil UUID)
2. Create a BEFORE INSERT/UPDATE trigger that auto-populates them from:
   - For data rows: extract from JSONB value->'__data__'->'metadata'
   - For all rows: fall back to app.tenant_id / app.node_id session variables
3. Enable RLS using the dedicated columns (same pattern as other tables)

Revision ID: 0b1c2d3e4f5a
Revises: 9a3b4c5d6e7f
Create Date: 2026-03-30 10:00:00.000000

"""
from typing import Sequence, Union

from alembic import op

# revision identifiers, used by Alembic.
revision: str = "0b1c2d3e4f5a"
down_revision: Union[str, None] = "9a3b4c5d6e7f"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

TABLE = "data_hierarchical_nodes"


def upgrade() -> None:
    op.execute(f"""
        DO $$
        BEGIN
            -- Pre-create the table matching LlamaIndex PostgresDocumentStore schema
            CREATE TABLE IF NOT EXISTS {TABLE} (
                id SERIAL NOT NULL,
                key VARCHAR NOT NULL,
                namespace VARCHAR NOT NULL,
                value JSONB,
                PRIMARY KEY (id),
                CONSTRAINT "{TABLE}:unique_key_namespace" UNIQUE (key, namespace)
            );
            CREATE INDEX IF NOT EXISTS "{TABLE}:idx_key_namespace"
                ON {TABLE} (key, namespace);

            -- 1. Add tenant_id and node_id columns
            IF NOT EXISTS (
                SELECT 1 FROM information_schema.columns
                WHERE table_name = '{TABLE}' AND column_name = 'tenant_id'
            ) THEN
                ALTER TABLE {TABLE}
                    ADD COLUMN tenant_id UUID NOT NULL DEFAULT '00000000-0000-0000-0000-000000000000'::uuid,
                    ADD COLUMN node_id UUID NOT NULL DEFAULT '00000000-0000-0000-0000-000000000000'::uuid;
            END IF;

            -- 2. Backfill existing data rows from JSONB metadata
            UPDATE {TABLE} SET
                tenant_id = COALESCE(
                    (value->'__data__'->'metadata'->>'tenant_id')::uuid,
                    '00000000-0000-0000-0000-000000000000'::uuid
                ),
                node_id = COALESCE(
                    (value->'__data__'->'metadata'->>'node_id')::uuid,
                    '00000000-0000-0000-0000-000000000000'::uuid
                )
            WHERE value ? '__data__';

            -- 3. Create trigger function to auto-populate tenant_id/node_id on INSERT/UPDATE
            CREATE OR REPLACE FUNCTION set_hierarchical_node_tenant() RETURNS trigger AS $fn$
            BEGIN
                -- Try extracting from JSONB metadata (data rows)
                IF NEW.value ? '__data__' THEN
                    NEW.tenant_id := COALESCE(
                        (NEW.value->'__data__'->'metadata'->>'tenant_id')::uuid,
                        NULLIF(current_setting('app.tenant_id', true), '')::uuid,
                        '00000000-0000-0000-0000-000000000000'::uuid
                    );
                    NEW.node_id := COALESCE(
                        (NEW.value->'__data__'->'metadata'->>'node_id')::uuid,
                        NULLIF(current_setting('app.node_id', true), '')::uuid,
                        '00000000-0000-0000-0000-000000000000'::uuid
                    );
                ELSE
                    -- Metadata rows: use session variables
                    NEW.tenant_id := COALESCE(
                        NULLIF(current_setting('app.tenant_id', true), '')::uuid,
                        '00000000-0000-0000-0000-000000000000'::uuid
                    );
                    NEW.node_id := COALESCE(
                        NULLIF(current_setting('app.node_id', true), '')::uuid,
                        '00000000-0000-0000-0000-000000000000'::uuid
                    );
                END IF;
                RETURN NEW;
            END;
            $fn$ LANGUAGE plpgsql;

            -- 4. Attach trigger
            DROP TRIGGER IF EXISTS trg_set_hierarchical_node_tenant ON {TABLE};
            CREATE TRIGGER trg_set_hierarchical_node_tenant
                BEFORE INSERT OR UPDATE ON {TABLE}
                FOR EACH ROW EXECUTE FUNCTION set_hierarchical_node_tenant();

            -- 5. Enable RLS with dedicated columns (same pattern as other tables)
            ALTER TABLE {TABLE} ENABLE ROW LEVEL SECURITY;
            ALTER TABLE {TABLE} FORCE ROW LEVEL SECURITY;

            DROP POLICY IF EXISTS tenant_isolation_{TABLE} ON {TABLE};
            CREATE POLICY tenant_isolation_{TABLE} ON {TABLE}
            FOR ALL
            USING (
                tenant_id = COALESCE(
                    NULLIF(current_setting('app.tenant_id', true), '')::uuid,
                    '00000000-0000-0000-0000-000000000000'::uuid
                )
            )
            WITH CHECK (
                tenant_id = COALESCE(
                    NULLIF(current_setting('app.tenant_id', true), '')::uuid,
                    '00000000-0000-0000-0000-000000000000'::uuid
                )
            );

            -- 6. Index for RLS performance
            CREATE INDEX IF NOT EXISTS idx_{TABLE}_tenant_id ON {TABLE} (tenant_id);

        END $$;
    """)


def downgrade() -> None:
    op.execute(f"""
        DO $$
        BEGIN
            DROP POLICY IF EXISTS tenant_isolation_{TABLE} ON {TABLE};
            ALTER TABLE {TABLE} DISABLE ROW LEVEL SECURITY;
            DROP TRIGGER IF EXISTS trg_set_hierarchical_node_tenant ON {TABLE};
            DROP FUNCTION IF EXISTS set_hierarchical_node_tenant();
            DROP INDEX IF EXISTS idx_{TABLE}_tenant_id;
            ALTER TABLE {TABLE} DROP COLUMN IF EXISTS tenant_id;
            ALTER TABLE {TABLE} DROP COLUMN IF EXISTS node_id;
        END $$;
    """)