"""add new column MetadataDescriptor.data_type + migrate data

Revision ID: 9e5b3a9b0cee
Revises: 8fbd40afbe5f
Create Date: 2023-03-22 13:26:27.580633

"""

from alembic import op
import sqlalchemy as sa
import sqlalchemy_utils
from sqlalchemy import orm
from sqlalchemy.ext.declarative import declarative_base
from zou.migrations.utils.base import BaseMixin
from sqlalchemy.dialects import postgresql


# revision identifiers, used by Alembic.
revision = "9e5b3a9b0cee"
down_revision = "8fbd40afbe5f"
branch_labels = None
depends_on = None


METADATA_DESCRIPTOR_TYPES = [
    ("string", "String"),
    ("number", "Number"),
    ("list", "List"),
    ("boolean", "Boolean"),
    ("checklist", "Checklist"),
]

base = declarative_base()


class Department(base, BaseMixin):
    """
    Studio department like modeling, animation, etc.
    """

    __tablename__ = "department"
    name = sa.Column(sa.String(80), unique=True, nullable=False)
    color = sa.Column(sa.String(7), nullable=False)


department_metadata_descriptor_link = sa.Table(
    "department_metadata_descriptor_link",
    base.metadata,
    sa.Column(
        "metadata_descriptor_id",
        sqlalchemy_utils.UUIDType(binary=False),
        sa.ForeignKey("metadata_descriptor.id"),
    ),
    sa.Column(
        "department_id",
        sqlalchemy_utils.UUIDType(binary=False),
        sa.ForeignKey("department.id"),
    ),
)


class MetadataDescriptor(base, BaseMixin):
    """
    This models allow to identify which metadata are available for a given
    project and a given entity type.
    """

    __tablename__ = "metadata_descriptor"

    project_id = sa.Column(
        sqlalchemy_utils.UUIDType(binary=False),
        sa.ForeignKey("project.id"),
        nullable=False,
        index=True,
    )
    entity_type = sa.Column(sa.String(60), nullable=False, index=True)
    name = sa.Column(sa.String(120), nullable=False)
    data_type = sa.Column(
        sqlalchemy_utils.types.choice.ChoiceType(METADATA_DESCRIPTOR_TYPES)
    )
    field_name = sa.Column(sa.String(120), nullable=False)
    choices = sa.Column(postgresql.JSONB(astext_type=sa.Text()))
    for_client = sa.Column(sa.Boolean(), default=False, index=True)
    departments = orm.relationship(
        "Department", secondary=department_metadata_descriptor_link
    )

    __table_args__ = (
        sa.UniqueConstraint(
            "project_id", "entity_type", "name", name="metadata_descriptor_uc"
        ),
        sa.UniqueConstraint(
            "project_id",
            "entity_type",
            "field_name",
            name="metadata_descriptor_uc2",
        ),
    )


PROJECT_STYLES = [
    ("2d", "2D Animation"),
    ("3d", "3D Animation"),
    ("2d3d", "2D/3D Animation"),
    ("vfx", "VFX"),
    ("stop-motion", "Stop Motion"),
    ("motion-design", "Motion Design"),
    ("archviz", "Archviz"),
    ("commercial", "Commercial"),
    ("catalog", "Catalog"),
]


class Project(base, BaseMixin):
    __tablename__ = "project"

    name = sa.Column(sa.String(80), nullable=False, unique=True, index=True)
    code = sa.Column(sa.String(80))
    description = sa.Column(sa.Text())
    shotgun_id = sa.Column(sa.Integer)
    file_tree = sa.Column(postgresql.JSONB(astext_type=sa.Text()))
    data = sa.Column(postgresql.JSONB(astext_type=sa.Text()))
    has_avatar = sa.Column(sa.Boolean(), default=False)
    fps = sa.Column(sa.String(10))
    ratio = sa.Column(sa.String(10))
    resolution = sa.Column(sa.String(12))
    production_type = sa.Column(sa.String(20), default="short")
    production_style = sa.Column(
        sqlalchemy_utils.types.choice.ChoiceType(PROJECT_STYLES)
    )
    start_date = sa.Column(sa.Date())
    end_date = sa.Column(sa.Date())
    man_days = sa.Column(sa.Integer)
    nb_episodes = sa.Column(sa.Integer, default=0)
    episode_span = sa.Column(sa.Integer, default=0)
    max_retakes = sa.Column(sa.Integer, default=0)
    is_clients_isolated = sa.Column(sa.Boolean(), default=False)

    project_status_id = sa.Column(
        sqlalchemy_utils.UUIDType(binary=False),
        sa.ForeignKey("project_status.id"),
        index=True,
    )


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("metadata_descriptor", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column(
                "data_type",
                sqlalchemy_utils.types.choice.ChoiceType(
                    METADATA_DESCRIPTOR_TYPES
                ),
                nullable=True,
            )
        )
    # ### end Alembic commands ###

    bind = op.get_bind()
    session = orm.Session(bind=bind)
    metadata_descriptors = session.query(MetadataDescriptor).all()
    for metadata_descriptor in metadata_descriptors:
        if len(metadata_descriptor.choices) > 1 and all(
            isinstance(s, str) for s in metadata_descriptor.choices
        ):
            if all(
                s.startswith("[ ]") or s.startswith("[x]")
                for s in metadata_descriptor.choices
            ):
                metadata_descriptor.data_type = "checklist"
            else:
                metadata_descriptor.data_type = "list"
        else:
            metadata_descriptor.data_type = "string"
    session.commit()


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("metadata_descriptor", schema=None) as batch_op:
        batch_op.drop_column("data_type")

    # ### end Alembic commands ###
