"""
初始化迁移脚本
Revision ID: initial
Revises: 
Create Date: 2026-03-14
"""
from alembic import op
import sqlalchemy as sa

# revision identifiers
revision = 'initial'
down_revision = None
branch_labels = None
depends_on = None


def upgrade():
    """创建表"""
    # 用户表
    op.create_table(
        'users',
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('username', sa.String(80), unique=True, nullable=False),
        sa.Column('password_hash', sa.String(256), nullable=False),
        sa.Column('email', sa.String(120), unique=True, nullable=False),
        sa.Column('nickname', sa.String(80)),
        sa.Column('role', sa.String(20), default='user'),
        sa.Column('status', sa.String(20), default='active'),
        sa.Column('created_at', sa.DateTime, default=sa.func.now()),
        sa.Column('last_login_at', sa.DateTime),
    )
    
    # Gateway 表
    op.create_table(
        'gateways',
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('name', sa.String(80), unique=True, nullable=False),
        sa.Column('url', sa.String(256), nullable=False),
        sa.Column('token_hash', sa.String(256)),
        sa.Column('status', sa.String(20), default='offline'),
        sa.Column('agent_count', sa.Integer, default=0),
        sa.Column('connection_limit', sa.Integer, default=10),
        sa.Column('heartbeat_interval', sa.Integer, default=60),
        sa.Column('allowed_ips', sa.JSON),
        sa.Column('last_heartbeat', sa.DateTime),
        sa.Column('created_at', sa.DateTime, default=sa.func.now()),
    )
    
    # Agent 表
    op.create_table(
        'agents',
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('name', sa.String(80), nullable=False),
        sa.Column('display_name', sa.String(80)),
        sa.Column('gateway_id', sa.String(36), sa.ForeignKey('gateways.id')),
        sa.Column('socket_id', sa.String(100)),
        sa.Column('model', sa.String(80)),
        sa.Column('capabilities', sa.JSON),
        sa.Column('status', sa.String(20), default='offline'),
        sa.Column('priority', sa.Integer, default=5),
        sa.Column('weight', sa.Integer, default=10),
        sa.Column('connection_limit', sa.Integer, default=5),
        sa.Column('current_sessions', sa.Integer, default=0),
        sa.Column('last_heartbeat', sa.DateTime),
        sa.Column('created_at', sa.DateTime, default=sa.func.now()),
    )
    
    # 会话表
    op.create_table(
        'sessions',
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('user_id', sa.String(36), sa.ForeignKey('users.id'), nullable=False),
        sa.Column('primary_agent_id', sa.String(36), sa.ForeignKey('agents.id')),
        sa.Column('participating_agent_ids', sa.JSON),
        sa.Column('user_socket_id', sa.String(100)),
        sa.Column('title', sa.String(200)),
        sa.Column('channel_type', sa.String(20), default='web'),
        sa.Column('status', sa.String(20), default='active'),
        sa.Column('message_count', sa.Integer, default=0),
        sa.Column('unread_count', sa.Integer, default=0),
        sa.Column('created_at', sa.DateTime, default=sa.func.now()),
        sa.Column('updated_at', sa.DateTime),
        sa.Column('last_active_at', sa.DateTime),
    )
    
    # 消息表
    op.create_table(
        'messages',
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('session_id', sa.String(36), sa.ForeignKey('sessions.id'), nullable=False),
        sa.Column('sender_type', sa.String(20), nullable=False),
        sa.Column('sender_id', sa.String(36), nullable=False),
        sa.Column('message_type', sa.String(20), default='text'),
        sa.Column('content', sa.Text),
        sa.Column('content_type', sa.String(20), default='markdown'),
        sa.Column('reply_to', sa.String(36)),
        sa.Column('status', sa.String(20), default='sent'),
        sa.Column('ack_status', sa.String(20), default='pending'),
        sa.Column('retry_count', sa.Integer, default=0),
        sa.Column('created_at', sa.DateTime, default=sa.func.now()),
        sa.Column('delivered_at', sa.DateTime),
    )
    
    # 连接表
    op.create_table(
        'connections',
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('socket_id', sa.String(100), unique=True, nullable=False),
        sa.Column('connection_type', sa.String(20), nullable=False),
        sa.Column('entity_id', sa.String(36), nullable=False),
        sa.Column('entity_type', sa.String(20), nullable=False),
        sa.Column('ip_address', sa.String(45)),
        sa.Column('user_agent', sa.String(500)),
        sa.Column('status', sa.String(20), default='connected'),
        sa.Column('auth_token', sa.String(500)),
        sa.Column('connected_at', sa.DateTime, default=sa.func.now()),
        sa.Column('last_activity', sa.DateTime),
        sa.Column('disconnected_at', sa.DateTime),
    )
    
    # 创建索引
    op.create_index('ix_messages_session_id', 'messages', ['session_id'])
    op.create_index('ix_sessions_user_id', 'sessions', ['user_id'])
    op.create_index('ix_sessions_status', 'sessions', ['status'])


def downgrade():
    """删除表"""
    op.drop_table('connections')
    op.drop_table('messages')
    op.drop_table('sessions')
    op.drop_table('agents')
    op.drop_table('gateways')
    op.drop_table('users')