"""Phase 2: Backfill existing data with pilot tenant/node UUIDs.

Data isolation spec Section 2.2 Phase 2.
Run AFTER Phase 1 migration.

Pilot UUIDs are sourced from environment variables PILOT_TENANT_ID and
PILOT_NODE_ID. Falls back to well-known pilot defaults if not set.

Revision ID: 8a2b3c4d5e6f
Revises: 7a1b2c3d4e5f
Create Date: 2026-03-13 10:00:01.000000

"""
import os
from typing import Sequence, Union

from alembic import op

# revision identifiers, used by Alembic.
revision: str = "8a2b3c4d5e6f"
down_revision: Union[str, None] = "7a1b2c3d4e5f"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

# Pilot UUIDs — override via env vars for your deployment
PILOT_TENANT_ID = os.environ.get(
    "PILOT_TENANT_ID", "00000000-0000-0000-0000-000000000001"
)
PILOT_NODE_ID = os.environ.get(
    "PILOT_NODE_ID", "00000000-0000-0000-0000-000000000001"
)


def upgrade() -> None:
    # Backfill documents
    op.execute(
        f"""
        UPDATE documents
        SET tenant_id = '{PILOT_TENANT_ID}', node_id = '{PILOT_NODE_ID}'
        WHERE tenant_id IS NULL
        """
    )

    # Backfill extracted_records (inherit from parent document)
    op.execute(
        f"""
        UPDATE extracted_records er
        SET tenant_id = d.tenant_id, node_id = d.node_id
        FROM documents d
        WHERE er.document_id = d.id AND er.tenant_id IS NULL
        """
    )

    # Backfill domain_term_dictionary (inherit from source document)
    op.execute(
        f"""
        UPDATE domain_term_dictionary dt
        SET tenant_id = d.tenant_id, node_id = d.node_id
        FROM documents d
        WHERE dt.source_document_id = d.id AND dt.tenant_id IS NULL
        """
    )

    # Backfill orphaned domain terms (no source_document_id)
    op.execute(
        f"""
        UPDATE domain_term_dictionary
        SET tenant_id = '{PILOT_TENANT_ID}', node_id = '{PILOT_NODE_ID}'
        WHERE tenant_id IS NULL
        """
    )

    # Backfill symbol_dictionaries (inherit from parent document)
    op.execute(
        f"""
        UPDATE symbol_dictionaries sd
        SET tenant_id = d.tenant_id, node_id = d.node_id
        FROM documents d
        WHERE sd.document_id = d.id AND sd.tenant_id IS NULL
        """
    )


def downgrade() -> None:
    # Reset all tenant/node columns to NULL
    for table in [
        "documents",
        "extracted_records",
        "domain_term_dictionary",
        "symbol_dictionaries",
    ]:
        op.execute(f"UPDATE {table} SET tenant_id = NULL, node_id = NULL")