"""Add workspace.planning_mode and child.preschool_planning_days_per_year.

Revision ID: 002
Revises: 001
Create Date: 2026-09-12
"""

from __future__ import annotations

from alembic import op

revision: str = "002"
down_revision: str | None = "001"
branch_labels: str | None = None
depends_on: str | None = None

_UPGRADE_SQL = """
ALTER TABLE workspace ADD COLUMN planning_mode TEXT NOT NULL DEFAULT 'solo'
    CHECK (planning_mode IN ('solo','couple'));

ALTER TABLE child ADD COLUMN preschool_planning_days_per_year INTEGER NOT NULL DEFAULT 6
    CHECK (preschool_planning_days_per_year >= 0 AND preschool_planning_days_per_year <= 30);
"""

_DOWNGRADE_SQL = """
ALTER TABLE child DROP COLUMN preschool_planning_days_per_year;
ALTER TABLE workspace DROP COLUMN planning_mode;
"""


def upgrade() -> None:
    from sqlalchemy import text

    for stmt in [s.strip() for s in _UPGRADE_SQL.split(";") if s.strip()]:
        op.execute(text(stmt))


def downgrade() -> None:
    from sqlalchemy import text

    for stmt in [s.strip() for s in _DOWNGRADE_SQL.split(";") if s.strip()]:
        op.execute(text(stmt))