"""add universe_snapshot tables

Revision ID: e6f7a8b9c0d1
Revises: d5e6f7a8b9c0
Create Date: 2026-03-29
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

revision: str = "e6f7a8b9c0d1"
down_revision: Union[str, Sequence[str], None] = "d5e6f7a8b9c0"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


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

    if not conn.dialect.has_table(conn, "universe_ticker_registry"):
        op.create_table(
            "universe_ticker_registry",
            sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
            sa.Column("ticker", sa.String(10), nullable=False),
            sa.Column("name", sa.String(255), nullable=True),
            sa.Column("cik", sa.String(20), nullable=True),
            sa.Column("sector", sa.String(100), nullable=True),
            sa.Column("industry", sa.String(200), nullable=True),
            sa.Column("exchange", sa.String(20), nullable=True),
            sa.Column("is_active", sa.Boolean(), nullable=False, server_default="true"),
            sa.Column("created_at", postgresql.TIMESTAMP(timezone=True)),
            sa.Column("updated_at", postgresql.TIMESTAMP(timezone=True)),
            sa.UniqueConstraint("ticker", name="uq_universe_ticker_registry"),
        )
        op.create_index("idx_registry_ticker", "universe_ticker_registry", ["ticker"])
        op.create_index("idx_registry_sector", "universe_ticker_registry", ["sector"])
        op.create_index("idx_registry_exchange", "universe_ticker_registry", ["exchange"])
        op.create_index("idx_registry_active", "universe_ticker_registry", ["is_active"])

    if not conn.dialect.has_table(conn, "universe_snapshot"):
        op.create_table(
            "universe_snapshot",
            sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
            sa.Column("ticker", sa.String(10), nullable=False),
            sa.Column("snapshot_date", postgresql.TIMESTAMP(timezone=True), nullable=False),
            sa.Column("close_price", sa.Float(), nullable=True),
            sa.Column("shares_outstanding", sa.Float(), nullable=True),
            sa.Column("market_cap", sa.Float(), nullable=True),
            sa.Column("sector", sa.String(100), nullable=True),
            sa.Column("industry", sa.String(200), nullable=True),
            sa.Column("exchange", sa.String(20), nullable=True),
            sa.Column("created_at", postgresql.TIMESTAMP(timezone=True)),
            sa.UniqueConstraint("ticker", "snapshot_date", name="uq_universe_snapshot"),
        )
        op.create_index("idx_snapshot_ticker_date", "universe_snapshot", ["ticker", "snapshot_date"])
        op.create_index("idx_snapshot_date_mcap", "universe_snapshot", ["snapshot_date", "market_cap"])
        op.create_index("idx_snapshot_sector", "universe_snapshot", ["sector"])
        op.create_index("idx_snapshot_date", "universe_snapshot", ["snapshot_date"])


def downgrade() -> None:
    op.drop_table("universe_snapshot")
    op.drop_table("universe_ticker_registry")