"""backfill project info in setup drafts

Revision ID: 20260511_01
Revises: 20260509_03
Create Date: 2026-05-11 10:30:00.000000

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


revision: str = "20260511_01"
down_revision: Union[str, None] = "20260509_03"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


EMPTY_PROJECT_INFO_JSON = """{
  "code": "",
  "name": "",
  "project_full_name": "",
  "sponsor": "",
  "protocol_no": "",
  "lead_unit": "",
  "principal_investigator": "",
  "main_pm": "",
  "research_analysis": "",
  "research_product": "",
  "control_product": "",
  "indication": "",
  "research_population": "",
  "research_design": "",
  "plan_start_date": "",
  "plan_end_date": "",
  "planned_site_count": null,
  "planned_enrollment_count": null,
  "status": "",
  "visit_schedule": []
}"""


def _has_table(inspector: sa.Inspector, table_name: str) -> bool:
    return table_name in set(inspector.get_table_names())


def _has_columns(inspector: sa.Inspector, table_name: str, column_names: set[str]) -> bool:
    existing = {col["name"] for col in inspector.get_columns(table_name)}
    return column_names.issubset(existing)


def _backfill_config_column(table_name: str, config_column: str, *, require_existing_setup_content: bool) -> None:
    content_filter = ""
    if require_existing_setup_content:
        content_filter = f"""
              AND (
                jsonb_array_length(COALESCE(setup.{config_column}->'projectMilestones', '[]'::jsonb)) > 0
                OR COALESCE(setup.{config_column}->'enrollmentPlan', '{{}}'::jsonb) <> '{{"totalTarget": 0, "startDate": "", "endDate": "", "monthlyGoalNote": "", "stageBreakdown": ""}}'::jsonb
                OR jsonb_array_length(COALESCE(setup.{config_column}->'siteMilestones', '[]'::jsonb)) > 0
                OR jsonb_array_length(COALESCE(setup.{config_column}->'siteEnrollmentPlans', '[]'::jsonb)) > 0
                OR jsonb_array_length(COALESCE(setup.{config_column}->'monitoringStrategies', '[]'::jsonb)) > 0
                OR jsonb_array_length(COALESCE(setup.{config_column}->'centerConfirm', '[]'::jsonb)) > 0
              )
        """
    op.execute(
        sa.text(
            f"""
            UPDATE {table_name} AS setup
            SET {config_column} = jsonb_set(
              setup.{config_column},
              '{{projectInfo}}',
              jsonb_build_object(
                'code', COALESCE(studies.code, ''),
                'name', COALESCE(studies.name, ''),
                'project_full_name', COALESCE(studies.project_full_name, ''),
                'sponsor', COALESCE(studies.sponsor, ''),
                'protocol_no', COALESCE(studies.protocol_no, ''),
                'lead_unit', COALESCE(studies.lead_unit, ''),
                'principal_investigator', COALESCE(studies.principal_investigator, ''),
                'main_pm', COALESCE(studies.main_pm, ''),
                'research_analysis', COALESCE(studies.research_analysis, ''),
                'research_product', COALESCE(studies.research_product, ''),
                'control_product', COALESCE(studies.control_product, ''),
                'indication', COALESCE(studies.indication, ''),
                'research_population', COALESCE(studies.research_population, ''),
                'research_design', COALESCE(studies.research_design, ''),
                'plan_start_date', COALESCE(to_char(studies.plan_start_date, 'YYYY-MM-DD'), ''),
                'plan_end_date', COALESCE(to_char(studies.plan_end_date, 'YYYY-MM-DD'), ''),
                'planned_site_count', to_jsonb(studies.planned_site_count),
                'planned_enrollment_count', to_jsonb(studies.planned_enrollment_count),
                'status', COALESCE(studies.status, ''),
                'visit_schedule', COALESCE(studies.visit_schedule::jsonb, '[]'::jsonb)
              ),
              true
            )
            FROM studies
            WHERE setup.study_id = studies.id
              AND setup.{config_column} IS NOT NULL
              AND (
                NOT setup.{config_column} ? 'projectInfo'
                OR setup.{config_column}->'projectInfo' = CAST(:empty_project_info AS jsonb)
              )
              {content_filter}
            """
        ).bindparams(empty_project_info=EMPTY_PROJECT_INFO_JSON)
    )


def upgrade() -> None:
    bind = op.get_bind()
    inspector = sa.inspect(bind)
    required_study_columns = {
        "id",
        "code",
        "name",
        "project_full_name",
        "sponsor",
        "protocol_no",
        "lead_unit",
        "principal_investigator",
        "main_pm",
        "research_analysis",
        "research_product",
        "control_product",
        "indication",
        "research_population",
        "research_design",
        "plan_start_date",
        "plan_end_date",
        "planned_site_count",
        "planned_enrollment_count",
        "status",
        "visit_schedule",
    }
    if not _has_table(inspector, "studies") or not _has_columns(inspector, "studies", required_study_columns):
        return

    if _has_table(inspector, "study_setup_configs") and _has_columns(
        inspector, "study_setup_configs", {"study_id", "config"}
    ):
        _backfill_config_column("study_setup_configs", "config", require_existing_setup_content=True)

    if _has_table(inspector, "study_setup_configs") and _has_columns(
        inspector, "study_setup_configs", {"study_id", "published_config"}
    ):
        _backfill_config_column("study_setup_configs", "published_config", require_existing_setup_content=False)

    if _has_table(inspector, "study_setup_config_versions") and _has_columns(
        inspector, "study_setup_config_versions", {"study_id", "config"}
    ):
        _backfill_config_column("study_setup_config_versions", "config", require_existing_setup_content=False)


def downgrade() -> None:
    pass