"""Add temporal dataset metadata and durable dataset-version provenance."""

from alembic import op
import sqlalchemy as sa


revision = "202607140001"
down_revision = "202606120900"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.add_column("datasets", sa.Column("temporal_series_key", sa.String(length=255), nullable=True))
    op.add_column("datasets", sa.Column("observed_at", sa.DateTime(timezone=True), nullable=True))
    op.add_column("datasets", sa.Column("valid_from", sa.DateTime(timezone=True), nullable=True))
    op.add_column("datasets", sa.Column("valid_to", sa.DateTime(timezone=True), nullable=True))
    op.add_column("datasets", sa.Column("temporal_granularity", sa.String(length=32), nullable=True))
    op.add_column("datasets", sa.Column("source_version", sa.String(length=120), nullable=True))

    op.add_column("dataset_versions", sa.Column("source_version", sa.String(length=120), nullable=True))
    op.add_column("dataset_versions", sa.Column("observed_at", sa.DateTime(timezone=True), nullable=True))
    op.add_column("dataset_versions", sa.Column("valid_from", sa.DateTime(timezone=True), nullable=True))
    op.add_column("dataset_versions", sa.Column("valid_to", sa.DateTime(timezone=True), nullable=True))
    op.add_column("dataset_versions", sa.Column("checksum_sha256", sa.String(length=64), nullable=True))
    op.add_column("dataset_versions", sa.Column("source_metadata", sa.JSON(), nullable=True))
    op.add_column("dataset_versions", sa.Column("provenance_metadata", sa.JSON(), nullable=True))

    op.create_index(
        "ix_datasets_project_temporal_series_observed",
        "datasets",
        ["project_id", "temporal_series_key", "observed_at"],
    )
    op.create_index("ix_dataset_versions_dataset_version", "dataset_versions", ["dataset_id", "version"], unique=True)
    op.create_index(
        "ix_vector_features_dataset_source_feature",
        "vector_features",
        ["dataset_id", "source_feature_id"],
    )
    op.create_check_constraint(
        "ck_datasets_temporal_valid_range",
        "datasets",
        "valid_to IS NULL OR valid_from IS NULL OR valid_to >= valid_from",
    )
    op.create_check_constraint(
        "ck_dataset_versions_temporal_valid_range",
        "dataset_versions",
        "valid_to IS NULL OR valid_from IS NULL OR valid_to >= valid_from",
    )


def downgrade() -> None:
    op.drop_constraint("ck_dataset_versions_temporal_valid_range", "dataset_versions", type_="check")
    op.drop_constraint("ck_datasets_temporal_valid_range", "datasets", type_="check")
    op.drop_index("ix_vector_features_dataset_source_feature", table_name="vector_features")
    op.drop_index("ix_dataset_versions_dataset_version", table_name="dataset_versions")
    op.drop_index("ix_datasets_project_temporal_series_observed", table_name="datasets")

    op.drop_column("dataset_versions", "provenance_metadata")
    op.drop_column("dataset_versions", "source_metadata")
    op.drop_column("dataset_versions", "checksum_sha256")
    op.drop_column("dataset_versions", "valid_to")
    op.drop_column("dataset_versions", "valid_from")
    op.drop_column("dataset_versions", "observed_at")
    op.drop_column("dataset_versions", "source_version")

    op.drop_column("datasets", "source_version")
    op.drop_column("datasets", "temporal_granularity")
    op.drop_column("datasets", "valid_to")
    op.drop_column("datasets", "valid_from")
    op.drop_column("datasets", "observed_at")
    op.drop_column("datasets", "temporal_series_key")