"""Add document_id column to data_hierarchical_nodes with FK CASCADE to documents.

Adds a dedicated document_id UUID column, backfills from JSONB metadata
(3 namespace paths), extends the existing trigger to auto-populate on INSERT,
and adds FK -> documents.id with ON DELETE CASCADE.

Revision ID: 2b3c4d5e6f7a
Revises: 1a2b3c4d5e6f
Create Date: 2026-04-09 10:00:01.000000

"""
from typing import Sequence, Union

from alembic import op

# revision identifiers, used by Alembic.
revision: str = "2b3c4d5e6f7a"
down_revision: Union[str, None] = "1a2b3c4d5e6f"
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
            -- 1. Add document_id column (nullable, no default)
            IF NOT EXISTS (
                SELECT 1 FROM information_schema.columns
                WHERE table_name = '{TABLE}' AND column_name = 'document_id'
            ) THEN
                ALTER TABLE {TABLE} ADD COLUMN document_id UUID;
            END IF;

            -- 2. Backfill from JSONB metadata (3 namespace paths)
            UPDATE {TABLE} SET document_id = (value->'__data__'->'metadata'->>'document_id')::uuid
            WHERE namespace = 'hierarchical/data'
            AND value->'__data__'->'metadata'->>'document_id' IS NOT NULL;

            UPDATE {TABLE} SET document_id = (value->>'ref_doc_id')::uuid
            WHERE namespace = 'hierarchical/metadata'
            AND value->>'ref_doc_id' IS NOT NULL;

            UPDATE {TABLE} SET document_id = key::uuid
            WHERE namespace = 'hierarchical/ref_doc_info'
            AND document_id IS NULL;

            -- 3. NULL out orphans (document_id not in documents table)
            UPDATE {TABLE} SET document_id = NULL
            WHERE document_id IS NOT NULL
            AND document_id NOT IN (SELECT id FROM documents);

            -- 4. Add FK constraint with CASCADE delete
            ALTER TABLE {TABLE}
                ADD CONSTRAINT fk_hierarchical_nodes_document
                FOREIGN KEY (document_id) REFERENCES documents(id) ON DELETE CASCADE;

            -- 5. Index for FK lookups and cascade performance
            CREATE INDEX IF NOT EXISTS idx_{TABLE}_document_id ON {TABLE} (document_id);

            -- 6. Extend trigger to auto-populate document_id on INSERT/UPDATE
            CREATE OR REPLACE FUNCTION set_hierarchical_node_tenant() RETURNS trigger AS $fn$
            BEGIN
                -- Data rows: extract from JSONB metadata
                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
                    );
                    NEW.document_id := (NEW.value->'__data__'->'metadata'->>'document_id')::uuid;
                ELSE
                    -- Metadata/ref_doc_info rows: use session variables for tenant
                    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
                    );
                    -- Try ref_doc_id for metadata rows
                    NEW.document_id := (NEW.value->>'ref_doc_id')::uuid;
                END IF;

                -- For ref_doc_info rows: key is the document_id
                IF NEW.document_id IS NULL AND NEW.namespace LIKE '%/ref_doc_info' THEN
                    BEGIN
                        NEW.document_id := NEW.key::uuid;
                    EXCEPTION WHEN invalid_text_representation THEN
                        NULL; -- key is not a valid UUID, leave document_id NULL
                    END;
                END IF;

                RETURN NEW;
            END;
            $fn$ LANGUAGE plpgsql;

        END $$;
    """)


def downgrade() -> None:
    # Restore original trigger without document_id
    op.execute(f"""
        DO $$
        BEGIN
            CREATE OR REPLACE FUNCTION set_hierarchical_node_tenant() RETURNS trigger AS $fn$
            BEGIN
                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
                    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;

            DROP INDEX IF EXISTS idx_{TABLE}_document_id;
            ALTER TABLE {TABLE} DROP CONSTRAINT IF EXISTS fk_hierarchical_nodes_document;
            ALTER TABLE {TABLE} DROP COLUMN IF EXISTS document_id;
        END $$;
    """)
