"""Initial schema with audio_tracks table

Revision ID: 001
Revises:
Create Date: 2025-11-27

"""
from typing import Sequence, Union

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

# revision identifiers, used by Alembic.
revision: str = '001'
down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Create pgvector extension
    op.execute('CREATE EXTENSION IF NOT EXISTS vector')
    op.execute('CREATE EXTENSION IF NOT EXISTS "uuid-ossp"')

    # Create audio_tracks table
    op.create_table(
        'audio_tracks',
        sa.Column('id', postgresql.UUID(as_uuid=True), server_default=sa.text('gen_random_uuid()'), nullable=False),
        sa.Column('filepath', sa.String(), nullable=False),
        sa.Column('filename', sa.String(), nullable=False),
        sa.Column('duration_seconds', sa.Float(), nullable=True),
        sa.Column('file_size_bytes', sa.BigInteger(), nullable=True),
        sa.Column('format', sa.String(), nullable=True),
        sa.Column('analyzed_at', sa.DateTime(), nullable=False, server_default=sa.text('now()')),

        # Musical features
        sa.Column('tempo_bpm', sa.Float(), nullable=True),
        sa.Column('key', sa.String(), nullable=True),
        sa.Column('time_signature', sa.String(), nullable=True),
        sa.Column('energy', sa.Float(), nullable=True),
        sa.Column('danceability', sa.Float(), nullable=True),
        sa.Column('valence', sa.Float(), nullable=True),
        sa.Column('loudness_lufs', sa.Float(), nullable=True),
        sa.Column('spectral_centroid', sa.Float(), nullable=True),
        sa.Column('zero_crossing_rate', sa.Float(), nullable=True),

        # Genre classification
        sa.Column('genre_primary', sa.String(), nullable=True),
        sa.Column('genre_secondary', postgresql.ARRAY(sa.String()), nullable=True),
        sa.Column('genre_confidence', sa.Float(), nullable=True),

        # Mood classification
        sa.Column('mood_primary', sa.String(), nullable=True),
        sa.Column('mood_secondary', postgresql.ARRAY(sa.String()), nullable=True),
        sa.Column('mood_arousal', sa.Float(), nullable=True),
        sa.Column('mood_valence', sa.Float(), nullable=True),

        # Instruments
        sa.Column('instruments', postgresql.ARRAY(sa.String()), nullable=True),

        # Vocals
        sa.Column('has_vocals', sa.Boolean(), nullable=True),
        sa.Column('vocal_gender', sa.String(), nullable=True),

        # Embeddings
        sa.Column('embedding', Vector(512), nullable=True),
        sa.Column('embedding_model', sa.String(), nullable=True),

        # Metadata
        sa.Column('metadata', postgresql.JSON(astext_type=sa.Text()), nullable=True),

        sa.PrimaryKeyConstraint('id')
    )

    # Create indexes
    op.create_index('idx_filepath', 'audio_tracks', ['filepath'], unique=True)
    op.create_index('idx_genre_primary', 'audio_tracks', ['genre_primary'])
    op.create_index('idx_mood_primary', 'audio_tracks', ['mood_primary'])
    op.create_index('idx_tempo_bpm', 'audio_tracks', ['tempo_bpm'])

    # Create vector index for similarity search (IVFFlat)
    # Note: This requires some data in the table to train the index
    # For now, we'll create it later when we have embeddings
    # op.execute(
    #     "CREATE INDEX idx_embedding ON audio_tracks USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100)"
    # )


def downgrade() -> None:
    op.drop_index('idx_tempo_bpm', table_name='audio_tracks')
    op.drop_index('idx_mood_primary', table_name='audio_tracks')
    op.drop_index('idx_genre_primary', table_name='audio_tracks')
    op.drop_index('idx_filepath', table_name='audio_tracks')
    op.drop_table('audio_tracks')
    op.execute('DROP EXTENSION IF EXISTS vector')
