"""phase 3.1 aqmm quality tags on existing service tasks

Revision ID: 20260625_phase_3_1_aqmm_task_quality_tags
Revises: 20260624_phase_3_aqmm_engagement_level_rework
Create Date: 2026-06-25
"""
from __future__ import annotations

from alembic import op
import sqlalchemy as sa

revision = "20260625_phase_3_1_aqmm_task_quality_tags"
down_revision = "20260624_phase_3_aqmm_engagement_level_rework"
branch_labels = None
depends_on = None


AQMM_TEMPLATE_COLUMNS = [
    sa.Column("is_aqmm_task", sa.Boolean(), nullable=False, server_default=sa.false()),
    sa.Column("aqmm_mandatory", sa.Boolean(), nullable=False, server_default=sa.false()),
    sa.Column("aqmm_evidence_required", sa.Boolean(), nullable=False, server_default=sa.false()),
    sa.Column("aqmm_manager_review_required", sa.Boolean(), nullable=False, server_default=sa.false()),
    sa.Column("aqmm_partner_review_required", sa.Boolean(), nullable=False, server_default=sa.false()),
    sa.Column("aqmm_review_partner_required", sa.Boolean(), nullable=False, server_default=sa.false()),
    sa.Column("aqmm_blocks_final_release", sa.Boolean(), nullable=False, server_default=sa.false()),
    sa.Column("aqmm_reference", sa.String(length=120), nullable=True),
]

AQMM_INSTANCE_COLUMNS = AQMM_TEMPLATE_COLUMNS + [
    sa.Column("aqmm_status", sa.String(length=40), nullable=False, server_default="not_required"),
    sa.Column("aqmm_review_status", sa.String(length=40), nullable=False, server_default="not_required"),
    sa.Column("aqmm_completed_at_utc", sa.DateTime(timezone=True), nullable=True),
]


def _has_table(bind, table_name: str) -> bool:
    return sa.inspect(bind).has_table(table_name)


def _has_column(bind, table_name: str, column_name: str) -> bool:
    if not _has_table(bind, table_name):
        return False
    return column_name in {c["name"] for c in sa.inspect(bind).get_columns(table_name)}


def _add_column_once(bind, table: str, column: sa.Column) -> None:
    if not _has_column(bind, table, column.name):
        with op.batch_alter_table(table) as batch:
            batch.add_column(column)


def _drop_column_once(bind, table: str, column_name: str) -> None:
    if _has_column(bind, table, column_name):
        with op.batch_alter_table(table) as batch:
            batch.drop_column(column_name)


def upgrade() -> None:
    bind = op.get_bind()

    for table in ["service_default_task_templates", "firm_service_task_templates"]:
        if _has_table(bind, table):
            for column in AQMM_TEMPLATE_COLUMNS:
                _add_column_once(bind, table, column.copy())

    if _has_table(bind, "client_service_task_instances"):
        for column in AQMM_INSTANCE_COLUMNS:
            _add_column_once(bind, "client_service_task_instances", column.copy())

        # Backfill existing generated tasks from their firm task templates.
        op.execute("""
            UPDATE client_service_task_instances csti
            SET is_aqmm_task = COALESCE(fst.is_aqmm_task, FALSE),
                aqmm_mandatory = COALESCE(fst.aqmm_mandatory, FALSE),
                aqmm_evidence_required = COALESCE(fst.aqmm_evidence_required, FALSE),
                aqmm_manager_review_required = COALESCE(fst.aqmm_manager_review_required, FALSE),
                aqmm_partner_review_required = COALESCE(fst.aqmm_partner_review_required, FALSE),
                aqmm_review_partner_required = COALESCE(fst.aqmm_review_partner_required, FALSE),
                aqmm_blocks_final_release = COALESCE(fst.aqmm_blocks_final_release, FALSE),
                aqmm_reference = fst.aqmm_reference,
                aqmm_status = CASE WHEN COALESCE(fst.is_aqmm_task, FALSE) THEN 'pending' ELSE 'not_required' END,
                aqmm_review_status = CASE
                    WHEN COALESCE(fst.aqmm_manager_review_required, FALSE)
                      OR COALESCE(fst.aqmm_partner_review_required, FALSE)
                      OR COALESCE(fst.aqmm_review_partner_required, FALSE)
                    THEN 'pending_review'
                    ELSE 'not_required'
                END
            FROM firm_service_task_templates fst
            WHERE csti.firm_task_template_id = fst.id
        """)


def downgrade() -> None:
    bind = op.get_bind()
    for table in ["client_service_task_instances"]:
        if _has_table(bind, table):
            for col in [
                "aqmm_completed_at_utc",
                "aqmm_review_status",
                "aqmm_status",
                "aqmm_reference",
                "aqmm_blocks_final_release",
                "aqmm_review_partner_required",
                "aqmm_partner_review_required",
                "aqmm_manager_review_required",
                "aqmm_evidence_required",
                "aqmm_mandatory",
                "is_aqmm_task",
            ]:
                _drop_column_once(bind, table, col)

    for table in ["firm_service_task_templates", "service_default_task_templates"]:
        if _has_table(bind, table):
            for col in [
                "aqmm_reference",
                "aqmm_blocks_final_release",
                "aqmm_review_partner_required",
                "aqmm_partner_review_required",
                "aqmm_manager_review_required",
                "aqmm_evidence_required",
                "aqmm_mandatory",
                "is_aqmm_task",
            ]:
                _drop_column_once(bind, table, col)
