"""add_rfp_document_entity

Revision ID: b86ade02428a
Revises: 3bd69f03a016
Create Date: 2026-03-01 06:56:56.876038

Creates the rfp_document table and migrates rfpqna from rfp_id -> rfp_document_id.
For each existing rfp that has qna rows, a corresponding rfp_document is created
and the FK is backfilled before making rfp_document_id NOT NULL.
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision: str = 'b86ade02428a'
down_revision: Union[str, Sequence[str], None] = '3bd69f03a016'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Upgrade schema."""

    # 1. Create rfp_document table
    op.create_table('rfp_document',
        sa.Column('rfp_profile_id', sa.Integer(), nullable=False),
        sa.Column('document_id', sa.Integer(), nullable=True),
        sa.Column('doc_type', sa.Enum('RFP', 'CONTRACT', 'SECURITY_QUESTIONNAIRE', 'COMPLIANCE', 'OTHER', name='rfpdocumenttype', native_enum=False), nullable=False),
        sa.Column('status', sa.Enum('PENDING', 'PROCESSING', 'EXTRACTED', 'ANSWERING', 'COMPLETED', 'FAILED', name='rfpdocumentstatus', native_enum=False), nullable=False),
        sa.Column('s3_url', sa.String(length=2048), nullable=True),
        sa.Column('uploaded_by_contact_id', sa.Integer(), nullable=True),
        sa.Column('additional_data', postgresql.JSONB(astext_type=sa.Text()), nullable=True),
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
        sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
        sa.Column('deleted_at', sa.DateTime(timezone=True), nullable=True),
        sa.ForeignKeyConstraint(['document_id'], ['document.id'], name=op.f('rfp_document_document_id_fkey'), ondelete='SET NULL'),
        sa.ForeignKeyConstraint(['rfp_profile_id'], ['rfp_profile.id'], name=op.f('rfp_document_rfp_profile_id_fkey'), ondelete='CASCADE'),
        sa.ForeignKeyConstraint(['uploaded_by_contact_id'], ['contact.id'], name=op.f('rfp_document_uploaded_by_contact_id_fkey'), ondelete='SET NULL'),
        sa.PrimaryKeyConstraint('id', name=op.f('rfp_document_pkey')),
        sa.UniqueConstraint('document_id', name=op.f('rfp_document_document_id_key'))
    )
    op.create_index('ix_rfp_document_profile_status', 'rfp_document', ['rfp_profile_id', 'status'], unique=False)
    op.create_index(op.f('rfp_document_doc_type_idx'), 'rfp_document', ['doc_type'], unique=False)
    op.create_index(op.f('rfp_document_id_idx'), 'rfp_document', ['id'], unique=False)
    op.create_index(op.f('rfp_document_rfp_profile_id_idx'), 'rfp_document', ['rfp_profile_id'], unique=False)
    op.create_index(op.f('rfp_document_status_idx'), 'rfp_document', ['status'], unique=False)
    op.create_index(op.f('rfp_document_uploaded_by_contact_id_idx'), 'rfp_document', ['uploaded_by_contact_id'], unique=False)

    # 2. Add a temporary column to track which legacy rfp.id each rfp_document came from
    op.add_column('rfp_document', sa.Column('_legacy_rfp_id', sa.Integer(), nullable=True))

    # 3. Backfill: create an rfp_document for each rfp that has qna rows
    #    Status mapping is direct (both use uppercase enum names).
    #    rfp_profile_id defaults to 0 if NULL (shouldn't happen, but safe fallback).
    op.execute("""
        INSERT INTO rfp_document (
            rfp_profile_id, document_id, doc_type, status,
            s3_url, uploaded_by_contact_id, additional_data,
            _legacy_rfp_id
        )
        SELECT
            COALESCE(r.rfp_profile_id, 0),
            r.document_id,
            'RFP',
            CASE r.status
                WHEN 'COMPLETED'  THEN 'COMPLETED'
                WHEN 'FAILED'     THEN 'FAILED'
                WHEN 'PROCESSING' THEN 'PROCESSING'
                ELSE 'PENDING'
            END,
            r.s3_url,
            r.uploaded_by_contact_id,
            r.additional_data,
            r.id
        FROM rfp r
        WHERE EXISTS (SELECT 1 FROM rfpqna q WHERE q.rfp_id = r.id)
    """)

    # 4. Add rfp_document_id to rfpqna as NULLABLE first
    op.add_column('rfpqna', sa.Column('rfp_document_id', sa.Integer(), nullable=True))
    op.create_index(op.f('rfpqna_rfp_document_id_idx'), 'rfpqna', ['rfp_document_id'], unique=False)

    # 5. Backfill rfpqna.rfp_document_id using the temporary mapping column
    op.execute("""
        UPDATE rfpqna q
        SET rfp_document_id = rd.id
        FROM rfp_document rd
        WHERE rd._legacy_rfp_id = q.rfp_id
    """)

    # 6. Drop the temporary column — no longer needed
    op.drop_column('rfp_document', '_legacy_rfp_id')

    # 7. Make rfp_document_id NOT NULL now that all rows are backfilled
    op.alter_column('rfpqna', 'rfp_document_id', nullable=False)

    # 8. Make uuid NOT NULL
    op.alter_column('rfpqna', 'uuid',
               existing_type=sa.UUID(),
               nullable=False)

    # 9. Drop old rfp_id FK/index, add new FKs
    op.drop_index(op.f('rfpqna_rfp_id_idx'), table_name='rfpqna')
    op.drop_constraint(op.f('rfpqna_rfp_id_fkey'), 'rfpqna', type_='foreignkey')
    op.create_foreign_key(op.f('rfpqna_rfp_document_id_fkey'), 'rfpqna', 'rfp_document', ['rfp_document_id'], ['id'], ondelete='CASCADE')
    op.create_foreign_key(op.f('rfpqna_parent_question_id_fkey'), 'rfpqna', 'rfpqna', ['parent_question_id'], ['id'], ondelete='SET NULL')
    op.drop_column('rfpqna', 'rfp_id')


def downgrade() -> None:
    """Downgrade schema."""

    # 1. Restore rfp_id column (nullable first)
    op.add_column('rfpqna', sa.Column('rfp_id', sa.INTEGER(), autoincrement=False, nullable=True))

    # 2. Add temporary mapping column back for reverse lookup
    op.add_column('rfp_document', sa.Column('_legacy_rfp_id', sa.Integer(), nullable=True))

    # 3. Rebuild the mapping: match rfp_document back to rfp via document_id
    op.execute("""
        UPDATE rfp_document rd
        SET _legacy_rfp_id = r.id
        FROM rfp r
        WHERE r.document_id = rd.document_id
          AND rd.document_id IS NOT NULL
    """)

    # 4. Backfill rfp_id from the mapping
    op.execute("""
        UPDATE rfpqna q
        SET rfp_id = rd._legacy_rfp_id
        FROM rfp_document rd
        WHERE q.rfp_document_id = rd.id
          AND rd._legacy_rfp_id IS NOT NULL
    """)

    # 5. Drop temp column
    op.drop_column('rfp_document', '_legacy_rfp_id')

    # 6. Make rfp_id NOT NULL (rows without a match will block this — acceptable for downgrade)
    op.alter_column('rfpqna', 'rfp_id', nullable=False)

    # 7. Drop new FKs, restore old
    op.drop_constraint(op.f('rfpqna_parent_question_id_fkey'), 'rfpqna', type_='foreignkey')
    op.drop_constraint(op.f('rfpqna_rfp_document_id_fkey'), 'rfpqna', type_='foreignkey')
    op.create_foreign_key(op.f('rfpqna_rfp_id_fkey'), 'rfpqna', 'rfp', ['rfp_id'], ['id'], ondelete='CASCADE')
    op.drop_index(op.f('rfpqna_rfp_document_id_idx'), table_name='rfpqna')
    op.create_index(op.f('rfpqna_rfp_id_idx'), 'rfpqna', ['rfp_id'], unique=False)
    op.alter_column('rfpqna', 'uuid',
               existing_type=sa.UUID(),
               nullable=True)
    op.drop_column('rfpqna', 'rfp_document_id')

    # 8. Drop rfp_document table
    op.drop_index(op.f('rfp_document_uploaded_by_contact_id_idx'), table_name='rfp_document')
    op.drop_index(op.f('rfp_document_status_idx'), table_name='rfp_document')
    op.drop_index(op.f('rfp_document_rfp_profile_id_idx'), table_name='rfp_document')
    op.drop_index(op.f('rfp_document_id_idx'), table_name='rfp_document')
    op.drop_index(op.f('rfp_document_doc_type_idx'), table_name='rfp_document')
    op.drop_index('ix_rfp_document_profile_status', table_name='rfp_document')
    op.drop_table('rfp_document')
