"""Initial migration

Revision ID: 488d0bd5ab01
Revises: 
Create Date: 2026-04-22 18:06:46.629263

"""
from typing import Sequence, Union

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

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


def upgrade() -> None:
    """Upgrade schema."""
    # users table
    op.create_table(
        'users',
        sa.Column('id', postgresql.UUID(as_uuid=True), server_default=sa.text('gen_random_uuid()'), nullable=False),
        sa.Column('email', sa.String(255), nullable=False),
        sa.Column('password_hash', sa.String(255), nullable=False),
        sa.Column('name', sa.String(100), nullable=True),
        sa.Column('plan', sa.String(20), server_default='free', nullable=False),
        sa.Column('max_queries', sa.Integer(), server_default='5', nullable=False),
        sa.Column('is_active', sa.Boolean(), server_default='true', nullable=False),
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('NOW()'), nullable=False),
        sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('NOW()'), nullable=False),
        sa.PrimaryKeyConstraint('id'),
        sa.UniqueConstraint('email')
    )

    # queries table
    op.create_table(
        'queries',
        sa.Column('id', postgresql.UUID(as_uuid=True), server_default=sa.text('gen_random_uuid()'), nullable=False),
        sa.Column('user_id', postgresql.UUID(as_uuid=True), nullable=False),
        sa.Column('keyword', sa.String(200), nullable=False),
        sa.Column('target_brand', sa.String(100), nullable=False),
        sa.Column('brand_aliases', postgresql.JSONB(), server_default='[]', nullable=False),
        sa.Column('platforms', postgresql.JSONB(), server_default='["wenxin", "kimi"]', nullable=False),
        sa.Column('frequency', sa.String(20), server_default='weekly', nullable=False),
        sa.Column('status', sa.String(20), server_default='active', nullable=False),
        sa.Column('last_queried_at', sa.DateTime(timezone=True), nullable=True),
        sa.Column('next_query_at', sa.DateTime(timezone=True), nullable=True),
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('NOW()'), nullable=False),
        sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('NOW()'), nullable=False),
        sa.PrimaryKeyConstraint('id'),
        sa.ForeignKeyConstraint(['user_id'], ['users.id'], ondelete='CASCADE')
    )
    op.create_index('idx_queries_user_id', 'queries', ['user_id'])
    op.create_index('idx_queries_status', 'queries', ['status'])
    op.create_index('idx_queries_next_query_at', 'queries', ['next_query_at'])

    # citation_records table
    op.create_table(
        'citation_records',
        sa.Column('id', postgresql.UUID(as_uuid=True), server_default=sa.text('gen_random_uuid()'), nullable=False),
        sa.Column('query_id', postgresql.UUID(as_uuid=True), nullable=False),
        sa.Column('platform', sa.String(50), nullable=False),
        sa.Column('cited', sa.Boolean(), server_default='false', nullable=False),
        sa.Column('citation_position', sa.Integer(), nullable=True),
        sa.Column('citation_text', sa.Text(), nullable=True),
        sa.Column('competitor_brands', postgresql.JSONB(), server_default='[]', nullable=False),
        sa.Column('raw_response', sa.Text(), nullable=True),
        sa.Column('queried_at', sa.DateTime(timezone=True), server_default=sa.text('NOW()'), nullable=False),
        sa.PrimaryKeyConstraint('id'),
        sa.ForeignKeyConstraint(['query_id'], ['queries.id'], ondelete='CASCADE')
    )
    op.create_index('idx_citation_records_query_id', 'citation_records', ['query_id'])
    op.create_index('idx_citation_records_queried_at', 'citation_records', ['queried_at'])
    op.create_index('idx_citation_records_platform', 'citation_records', ['platform'])

    # query_tasks table
    op.create_table(
        'query_tasks',
        sa.Column('id', postgresql.UUID(as_uuid=True), server_default=sa.text('gen_random_uuid()'), nullable=False),
        sa.Column('query_id', postgresql.UUID(as_uuid=True), nullable=False),
        sa.Column('platform', sa.String(50), nullable=False),
        sa.Column('status', sa.String(20), server_default='pending', nullable=False),
        sa.Column('error_message', sa.Text(), nullable=True),
        sa.Column('scheduled_at', sa.DateTime(timezone=True), server_default=sa.text('NOW()'), nullable=False),
        sa.Column('started_at', sa.DateTime(timezone=True), nullable=True),
        sa.Column('completed_at', sa.DateTime(timezone=True), nullable=True),
        sa.PrimaryKeyConstraint('id'),
        sa.ForeignKeyConstraint(['query_id'], ['queries.id'], ondelete='CASCADE')
    )
    op.create_index('idx_query_tasks_status', 'query_tasks', ['status'])

    # subscriptions table
    op.create_table(
        'subscriptions',
        sa.Column('id', postgresql.UUID(as_uuid=True), server_default=sa.text('gen_random_uuid()'), nullable=False),
        sa.Column('user_id', postgresql.UUID(as_uuid=True), nullable=False),
        sa.Column('plan', sa.String(20), nullable=False),
        sa.Column('status', sa.String(20), server_default='active', nullable=False),
        sa.Column('start_date', sa.Date(), nullable=False),
        sa.Column('end_date', sa.Date(), nullable=False),
        sa.Column('amount', sa.Numeric(10, 2), nullable=True),
        sa.Column('payment_method', sa.String(50), nullable=True),
        sa.Column('payment_id', sa.String(255), nullable=True),
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('NOW()'), nullable=False),
        sa.PrimaryKeyConstraint('id'),
        sa.ForeignKeyConstraint(['user_id'], ['users.id'], ondelete='CASCADE')
    )


def downgrade() -> None:
    """Downgrade schema."""
    op.drop_table('subscriptions')
    op.drop_index('idx_query_tasks_status', table_name='query_tasks')
    op.drop_table('query_tasks')
    op.drop_index('idx_citation_records_platform', table_name='citation_records')
    op.drop_index('idx_citation_records_queried_at', table_name='citation_records')
    op.drop_index('idx_citation_records_query_id', table_name='citation_records')
    op.drop_table('citation_records')
    op.drop_index('idx_queries_next_query_at', table_name='queries')
    op.drop_index('idx_queries_status', table_name='queries')
    op.drop_index('idx_queries_user_id', table_name='queries')
    op.drop_table('queries')
    op.drop_table('users')