"""Step 4: Recreate indexes and cleanup (final step)

=== STATE TRANSITION ===
INITIAL:  Tables after Step 3 schema alterations:
          - SQLite: indexes lost during batch_alter_table operations
          - PostgreSQL/MySQL: indexes may or may not be preserved
          - schema_migrations table (old SQLite migration tracking)
RESULT:   Tables with all performance indexes recreated:
          - idx_agents_owner, idx_shares_agent, idx_shares_principal
          - idx_conversations_user, idx_conversations_user_status
          - idx_messages_* (6 indexes), idx_ma_* (4 indexes)
          - schema_migrations table removed (replaced by Alembic)
          - Idempotent: only creates indexes if they don't exist

DOWNGRADE: Recreates schema_migrations table only
           Keeps indexes (old plugin handles them with IF NOT EXISTS)
           Restores old plugin compatibility
========================

Revision ID: 43647ca7796c
Revises: 35647ca7796c
Create Date: 2025-11-25 17:00:00.000000
"""

import os
import sys
from typing import Sequence, Union

import sqlalchemy as sa

from alembic import op

# Add alembic folder to path for utils import
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils import (
    create_index,
    drop_table,
    get_db_schema,
    get_table_indexes,
    timed_step,
)
from utils import (
    migration_logger as logger,
)
from utils import (
    table_exists as _table_exists,
)

# revision identifiers, used by Alembic.
revision: str = "43647ca7796c"
down_revision: Union[str, Sequence[str], None] = "35647ca7796c"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Recreate indexes and cleanup legacy artifacts.

    IMPORTANT: Step 3's batch_alter_table operations cause SQLite to lose all indexes.
    This step recreates them for performance. For PostgreSQL/MySQL, indexes may have
    been preserved, so we only create if missing (idempotent).
    """
    with timed_step("Step 4/4: Recreate indexes and cleanup"):
        bind = op.get_bind()
        inspector = sa.inspect(bind)

        # Drop legacy schema_migrations table
        if _table_exists("schema_migrations"):
            logger.info("  Dropping legacy schema_migrations table")
            drop_table("schema_migrations")

        logger.info("  Recreating indexes (idempotent)...")

        # Recreate indexes for agents
        if _table_exists("agents"):
            indexes = set(get_table_indexes("agents", inspector))
            if op.f("idx_agents_owner") not in indexes:
                logger.info("    Creating idx_agents_owner")
                create_index(op.f("idx_agents_owner"), "agents", ["owner"], schema=get_db_schema())

        # Recreate indexes for agent_shares
        if _table_exists("agent_shares"):
            indexes = set(get_table_indexes("agent_shares", inspector))
            if op.f("idx_shares_agent") not in indexes:
                logger.info("    Creating idx_shares_agent")
                create_index(op.f("idx_shares_agent"), "agent_shares", ["agent_id"], schema=get_db_schema())
            if op.f("idx_shares_principal") not in indexes:
                logger.info("    Creating idx_shares_principal")
                create_index(op.f("idx_shares_principal"), "agent_shares", ["principal"], schema=get_db_schema())

        # Recreate indexes for conversations
        if _table_exists("conversations"):
            indexes = set(get_table_indexes("conversations", inspector))
            if op.f("idx_conversations_user") not in indexes:
                logger.info("    Creating idx_conversations_user")
                create_index(op.f("idx_conversations_user"), "conversations", ["user_id"], schema=get_db_schema())
            if op.f("idx_conversations_user_status") not in indexes:
                logger.info("    Creating idx_conversations_user_status")
                create_index(op.f("idx_conversations_user_status"), "conversations", ["user_id", "status"], schema=get_db_schema())

        # Recreate indexes for messages
        if _table_exists("messages"):
            indexes = set(get_table_indexes("messages", inspector))
            if op.f("idx_messages_conversation") not in indexes:
                logger.info("    Creating idx_messages_conversation")
                create_index(op.f("idx_messages_conversation"), "messages", ["conversation_id"], schema=get_db_schema())
            if op.f("idx_messages_created") not in indexes:
                logger.info("    Creating idx_messages_created")
                create_index(op.f("idx_messages_created"), "messages", ["created_at"], schema=get_db_schema())
            if op.f("idx_messages_role") not in indexes:
                logger.info("    Creating idx_messages_role")
                create_index(op.f("idx_messages_role"), "messages", ["role"], schema=get_db_schema())
            if op.f("idx_messages_conv_created") not in indexes:
                logger.info("    Creating idx_messages_conv_created")
                create_index(op.f("idx_messages_conv_created"), "messages", ["conversation_id", "created_at"], schema=get_db_schema())
            if op.f("idx_messages_role_created") not in indexes:
                logger.info("    Creating idx_messages_role_created")
                create_index(op.f("idx_messages_role_created"), "messages", ["role", "created_at"], schema=get_db_schema())
            # Conditional index for active messages (SQLite-specific)
            if op.f("idx_messages_conv_created_active") not in indexes:
                try:
                    logger.info("    Creating idx_messages_conv_created_active (conditional)")
                    create_index(
                        op.f("idx_messages_conv_created_active"),
                        "messages",
                        ["conversation_id", "created_at"],
                        sqlite_where=sa.text("status='active'"),
                        schema=get_db_schema(),
                    )
                except Exception as e:
                    logger.warning(f"    Could not create conditional index: {e}")

        # Recreate indexes for message_agents
        if _table_exists("message_agents"):
            indexes = set(get_table_indexes("message_agents", inspector))
            if op.f("idx_ma_message") not in indexes:
                logger.info("    Creating idx_ma_message")
                create_index(op.f("idx_ma_message"), "message_agents", ["message_id"], schema=get_db_schema())
            if op.f("idx_ma_agent") not in indexes:
                logger.info("    Creating idx_ma_agent")
                create_index(op.f("idx_ma_agent"), "message_agents", ["agent_id"], schema=get_db_schema())
            if op.f("idx_ma_selected") not in indexes:
                logger.info("    Creating idx_ma_selected")
                create_index(op.f("idx_ma_selected"), "message_agents", ["selected"], schema=get_db_schema())
            if op.f("idx_ma_used") not in indexes:
                logger.info("    Creating idx_ma_used")
                create_index(op.f("idx_ma_used"), "message_agents", ["used"], schema=get_db_schema())

    logger.info("=" * 60 + " MIGRATION COMPLETE " + "=" * 60)


def downgrade() -> None:
    """Downgrade Step 4 - Recreate schema_migrations table.

    This reverses Step 4 by:
    - Recreating schema_migrations table with old structure
    - Inserting baseline marker (version 3)

    NOTE: We keep the indexes because old plugin's CREATE INDEX IF NOT EXISTS
    will skip recreation, making startup faster. Indexes are fully compatible.

    After downgrade, old v1.1.x plugin can work with the database.
    """
    logger.info("[Downgrade Step 4] Recreating legacy schema_migrations table...")

    bind = op.get_bind()

    # Recreate schema_migrations table for old plugin compatibility
    if not _table_exists("schema_migrations"):
        logger.info("  Recreating schema_migrations table")
        op.create_table(
            "schema_migrations",
            sa.Column("version", sa.Integer(), nullable=False),
            sa.Column("name", sa.Text(), nullable=False),
            sa.Column(
                "applied_at", sa.Text(), nullable=True, server_default=sa.text("(strftime('%Y-%m-%dT%H:%M:%fZ','now'))")
            ),
            sa.Column("checksum", sa.Text(), nullable=True),
            sa.Column("script", sa.Text(), nullable=True),
            sa.Column("runtime_ms", sa.Integer(), nullable=True),
            sa.PrimaryKeyConstraint("version"),
        )
        # Insert marker for baseline migration (version 3 = initial schema with migrations 1-3)
        bind.execute(
            sa.text(
                "INSERT INTO schema_migrations (version, name) VALUES (3, 'baseline (restored from Alembic downgrade)')"
            )
        )

    logger.info("[Downgrade Step 4] Complete")
    logger.info("[Note] Indexes kept - old plugin will recognize them with IF NOT EXISTS")
    logger.info("")
    logger.info("=" * 60)
    logger.info("DOWNGRADE COMPLETE")
    logger.info("=" * 60)
    logger.info("Your database is now compatible with v1.1.x plugin.")
    logger.info("You can safely downgrade the plugin version.")
    logger.info("=" * 60)
