"""Step 2: Create new tables (non-destructive)

=== STATE TRANSITION ===
INITIAL:  Old SQLite store base tables:
          - agents, conversations, messages, message_agents, preferences
          - draft_conversations, agent_shares, derived_documents (with 'metadata' column)
          - message_attachements (typo - old table name)
          All with TEXT columns, TEXT timestamps, no strict types
RESULT:   Same tables + NEW tables (admin_settings)
          + message_attachments (fixed typo from message_attachements)
          - Creates tables only if they don't exist (idempotent)
          - Old plugin ignores new tables

DOWNGRADE: Drops ONLY admin_settings (truly new in v1.2)
           Keeps all base schema tables including derived_documents
========================

Revision ID: 34647ca7796c
Revises: 33647ca7796c
Create Date: 2025-11-25 16: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 (
    JsonEncoded,
    create_table,
    drop_table,
    get_fk_prefix,
    get_prefixed_table_name,
    get_qualified_table_name,
    is_snowflake,
    timed_step,
)
from utils import (
    migration_logger as logger,
)
from utils import (
    table_exists as _table_exists,
)

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

STR_LEN = 255


def get_datetime_type():
    if op.get_bind().dialect.name == "mysql":
        from sqlalchemy.dialects.mysql import DATETIME

        return DATETIME(fsp=6)
    return sa.DateTime()


def get_current_timestamp():
    if op.get_bind().dialect.name == "mysql":
        return sa.text("CURRENT_TIMESTAMP(6)")
    return sa.func.current_timestamp()


def upgrade() -> None:
    """Create new tables required for v1.2.

    This step is NON-DESTRUCTIVE and BACKWARD COMPATIBLE.
    Old plugin versions will continue to work - they simply ignore new tables.
    """
    with timed_step("Step 2/4: Create new tables"):
        fk_prefix = get_fk_prefix()
        CREATED_AT = "created_at"
        LAST_MODIFIED = "last_modified"

        # Create agents table if it doesn't exist (fresh install)
        if not _table_exists("agents"):
            logger.info("  Creating agents table")
            create_table(
                "agents",
                sa.Column("id", sa.String(STR_LEN), nullable=False),
                sa.Column("owner", sa.String(STR_LEN), nullable=False),
                sa.Column("description", sa.Text(), nullable=True),
                sa.Column("sample_questions", JsonEncoded(), nullable=True),
                sa.Column("documents", JsonEncoded(), nullable=True),
                sa.Column("indexing", JsonEncoded(), nullable=True),
                sa.Column("published_version", JsonEncoded(), nullable=True),
                sa.Column("published_at", get_datetime_type(), nullable=True),
                sa.Column(
                    "publishing_status",
                    sa.Enum(
                        "idle", "publishing", "published", "failed", name="publishingstatusenum", native_enum=False
                    ),
                    nullable=True,
                    default="idle",
                ),
                sa.Column("publishing_job_id", sa.String(STR_LEN), nullable=True),
                sa.Column(CREATED_AT, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.Column(LAST_MODIFIED, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.PrimaryKeyConstraint("id"),
            )

        # Create agent_shares table if it doesn't exist
        if not _table_exists("agent_shares"):
            logger.info("  Creating agent_shares table")
            create_table(
                "agent_shares",
                sa.Column("agent_id", sa.String(STR_LEN), nullable=False),
                sa.Column("principal", sa.String(STR_LEN), nullable=False),
                sa.Column(
                    "principal_type",
                    sa.Enum("user", "group", name="principaltypeenum", native_enum=False),
                    nullable=False,
                ),
                sa.ForeignKeyConstraint(
                    ["agent_id"], [f"{fk_prefix}agents.id"], name=f"{fk_prefix}fk_agent_shares_agent_id_agents"
                ),
                sa.PrimaryKeyConstraint("agent_id", "principal"),
            )

        # Create conversations table if it doesn't exist
        if not _table_exists("conversations"):
            logger.info("  Creating conversations table")
            create_table(
                "conversations",
                sa.Column("conversation_id", sa.String(STR_LEN), nullable=False),
                sa.Column("user_id", sa.String(STR_LEN), nullable=False),
                sa.Column("title", sa.Text(), nullable=False),
                sa.Column("agent_ids", JsonEncoded(), nullable=False),
                sa.Column("llm_id", sa.String(STR_LEN), nullable=True),  # nullable for compatibility
                sa.Column("agents_enabled", sa.Boolean(), server_default=sa.sql.expression.true(), nullable=True),
                sa.Column(
                    "status",
                    sa.Enum("active", "deleted", name="statusenum", native_enum=False),
                    server_default=sa.text("'active'"),
                    nullable=False,
                ),
                sa.Column(CREATED_AT, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.Column(LAST_MODIFIED, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.PrimaryKeyConstraint("conversation_id"),
            )

        # Create draft_conversations table if it doesn't exist
        if not _table_exists("draft_conversations"):
            logger.info("  Creating draft_conversations table")
            create_table(
                "draft_conversations",
                sa.Column("agent_id", sa.String(STR_LEN), nullable=False),
                sa.Column("user_id", sa.String(STR_LEN), nullable=False),
                sa.Column("convo", JsonEncoded(), nullable=False),
                sa.Column(CREATED_AT, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.Column(LAST_MODIFIED, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.PrimaryKeyConstraint("agent_id", "user_id"),
            )

        # Create messages table if it doesn't exist
        if not _table_exists("messages"):
            logger.info("  Creating messages table")
            create_table(
                "messages",
                sa.Column("id", sa.String(STR_LEN), nullable=False),
                sa.Column("conversation_id", sa.String(STR_LEN), nullable=False),
                sa.Column("role", sa.String(STR_LEN), nullable=False),
                sa.Column("content", sa.Text(), nullable=False),
                sa.Column("event_log", JsonEncoded(), nullable=False),
                sa.Column(CREATED_AT, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.Column(LAST_MODIFIED, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.Column("actions", JsonEncoded(), nullable=False),
                sa.Column("artifacts", sa.LargeBinary(), nullable=True),
                sa.Column("selected_agent_ids", JsonEncoded(), nullable=False),
                sa.Column("used_agent_ids", JsonEncoded(), nullable=False),
                sa.Column("llm_id", sa.String(STR_LEN), nullable=True),  # nullable for compatibility
                sa.Column("agents_enabled", sa.Boolean(), server_default=sa.sql.expression.true(), nullable=True),
                sa.Column("feedback_rating", sa.Integer(), nullable=True),
                sa.Column("feedback_text", sa.Text(), nullable=True),
                sa.Column("feedback_by", sa.String(STR_LEN), nullable=True),
                sa.Column(
                    "feedback_updated_at", get_datetime_type(), nullable=True, server_default=get_current_timestamp()
                ),
                sa.Column("trace", sa.LargeBinary(), nullable=True),
                sa.Column(
                    "status",
                    sa.Enum("active", "deleted", name="statusenum", native_enum=False),
                    server_default=sa.text("'active'"),
                    nullable=False,
                ),
                sa.ForeignKeyConstraint(
                    ["conversation_id"],
                    [f"{fk_prefix}conversations.conversation_id"],
                    name=f"{fk_prefix}fk_messages_conversation_id_conversations",
                ),
                sa.PrimaryKeyConstraint("id"),
            )

        # Create message_agents table if it doesn't exist
        if not _table_exists("message_agents"):
            logger.info("  Creating message_agents table")
            table_args = [
                "message_agents",
                sa.Column("message_id", sa.String(STR_LEN), nullable=False),
                sa.Column("agent_id", sa.String(STR_LEN), nullable=False),
                sa.Column("selected", sa.Integer(), nullable=False),
                sa.Column("used", sa.Integer(), nullable=False),
            ]
            if not is_snowflake():
                table_args.extend(
                    [
                        sa.CheckConstraint("selected IN (0, 1)", name=f"{fk_prefix}ck_message_agents_selected"),
                        sa.CheckConstraint("used IN (0, 1)", name=f"{fk_prefix}ck_message_agents_used"),
                    ]
                )
            table_args.extend(
                [
                    sa.ForeignKeyConstraint(["message_id"], [f"{fk_prefix}messages.id"], name= f"{fk_prefix}fk_message_agents_message_id_messages" ,ondelete="CASCADE"),
                    sa.PrimaryKeyConstraint("message_id", "agent_id"),
                ]
            )
            create_table(*table_args)

        # Create preferences table if it doesn't exist
        if not _table_exists("preferences"):
            logger.info("  Creating preferences table")
            create_table(
                "preferences",
                sa.Column("user_id", sa.String(STR_LEN), nullable=False),
                sa.Column("prefs", JsonEncoded(), nullable=False),
                sa.Column(CREATED_AT, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.Column(LAST_MODIFIED, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.PrimaryKeyConstraint("user_id"),
            )

        # Create derived_documents table if it doesn't exist
        if not _table_exists("derived_documents"):
            logger.info("  Creating derived_documents table")
            create_table(
                "derived_documents",
                sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
                sa.Column("conv_id", sa.String(STR_LEN), nullable=False),
                sa.Column("user_id", sa.String(STR_LEN), nullable=False),
                sa.Column("document_path", sa.String(STR_LEN), nullable=False),
                sa.Column("document_name", sa.String(STR_LEN), nullable=False),
                sa.Column("document_metadata", JsonEncoded(), nullable=False),
                sa.Column(CREATED_AT, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.PrimaryKeyConstraint("id"),
                sa.UniqueConstraint("conv_id", "user_id", "document_path"),
            )

        # Create message_attachments table if it doesn't exist
        # AND migrate data from old typo table 'message_attachements' if it exists
        if not _table_exists("message_attachments"):
            logger.info("  Creating message_attachments table")
            create_table(
                "message_attachments",
                sa.Column("message_id", sa.String(STR_LEN), nullable=False),
                sa.Column("attachments", JsonEncoded(), nullable=False),
                sa.Column("extraction_mode", JsonEncoded(), nullable=True),
                sa.Column(CREATED_AT, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.Column(LAST_MODIFIED, get_datetime_type(), nullable=False, server_default=get_current_timestamp()),
                sa.PrimaryKeyConstraint("message_id"),
            )

            # Migrate data from old typo table 'message_attachements' if it exists
            if _table_exists("message_attachements"):
                logger.info("  Migrating data from message_attachements (typo) to message_attachments...")
                bind = op.get_bind()
                old_table_prefixed = get_prefixed_table_name("message_attachements")
                old_table = get_qualified_table_name("message_attachements")
                new_table = get_qualified_table_name("message_attachments")

                try:
                    # Check if extraction_mode column exists in old table
                    inspector = sa.inspect(bind)
                    old_columns = [col["name"] for col in inspector.get_columns(old_table_prefixed)]
                    has_extraction_mode = "extraction_mode" in old_columns

                    # Use INSERT...SELECT to copy data efficiently
                    # Map updated_at -> last_modified (old table used updated_at)
                    if has_extraction_mode:
                        bind.execute(
                            sa.text(
                                f"INSERT INTO {new_table} (message_id, attachments, extraction_mode, created_at, last_modified) "
                                f"SELECT message_id, attachments, extraction_mode, created_at, "
                                f"COALESCE(updated_at, created_at) "
                                f"FROM {old_table}"
                            )
                        )
                    else:
                        bind.execute(
                            sa.text(
                                f"INSERT INTO {new_table} (message_id, attachments, created_at, last_modified) "
                                f"SELECT message_id, attachments, created_at, "
                                f"COALESCE(updated_at, created_at) "
                                f"FROM {old_table}"
                            )
                        )
                    # Count migrated rows
                    count = bind.execute(sa.text(f"SELECT COUNT(*) FROM {new_table}")).scalar()
                    logger.info(f"    Successfully migrated {count} rows")
                except Exception as e:
                    logger.warning(f"    Could not migrate data from message_attachements: {e}")

        # Create admin_settings table if it doesn't exist
        if not _table_exists("admin_settings"):
            logger.info("  Creating admin_settings table")
            create_table(
                "admin_settings",
                sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
                sa.Column(CREATED_AT, sa.DateTime(), nullable=False, server_default=sa.func.current_timestamp()),
                sa.Column("user", sa.String(255), nullable=False),
                sa.Column("settings", JsonEncoded(), nullable=False),
                sa.PrimaryKeyConstraint("id"),
            )


def downgrade() -> None:
    """Drop ONLY new tables not in old base schema.

    IMPORTANT: Old SQLite store already has these base tables:
    - agents, conversations, messages, message_agents
    - preferences, draft_conversations, agent_shares
    - derived_documents (with 'metadata' column)
    - message_attachements (typo spelling)

    We ONLY drop tables that are truly NEW in v1.2:
    - admin_settings (new in v1.2)

    We do NOT drop:
    - derived_documents (existed in old schema with 'metadata' column)
    - message_attachments (new name, but old 'message_attachements' table exists)

    Dropping base tables would break old plugin compatibility!
    """
    logger.info("[Downgrade Step 2] Removing v1.2-only tables...")

    # Only drop tables that did NOT exist in old SQLite base schema
    if _table_exists("admin_settings"):
        logger.info("  Dropping admin_settings (new in v1.2)")
        drop_table("admin_settings")

    # Restore legacy typo table for old plugin compatibility
    if not _table_exists("message_attachements"):
        logger.info("  Recreating legacy message_attachements table (typo)")
        create_table(
            "message_attachements",
            sa.Column("message_id", sa.String(STR_LEN), nullable=False),
            sa.Column("attachments", JsonEncoded(), nullable=False),
            sa.Column("extraction_mode", JsonEncoded(), nullable=True),
            sa.Column("created_at", sa.Text(), nullable=True),
            sa.Column("updated_at", sa.Text(), nullable=True),
            sa.PrimaryKeyConstraint("message_id"),
        )

        if _table_exists("message_attachments"):
            logger.info("  Restoring data from message_attachments to message_attachements...")
            bind = op.get_bind()
            old_table_prefixed = get_prefixed_table_name("message_attachements")
            old_table = get_qualified_table_name("message_attachements")
            new_table = get_qualified_table_name("message_attachments")
            try:
                inspector = sa.inspect(bind)
                new_columns = [
                    col["name"] for col in inspector.get_columns(get_prefixed_table_name("message_attachments"))
                ]
                has_last_modified = "last_modified" in new_columns
                updated_at_source = "last_modified" if has_last_modified else "created_at"

                bind.execute(
                    sa.text(
                        f"INSERT INTO {old_table} (message_id, attachments, extraction_mode, created_at, updated_at) "
                        f"SELECT message_id, attachments, extraction_mode, created_at, "
                        f"COALESCE({updated_at_source}, created_at) "
                        f"FROM {new_table}"
                    )
                )
                count = bind.execute(sa.text(f"SELECT COUNT(*) FROM {old_table}")).scalar()
                logger.info(f"    Restored {count} rows")
            except Exception as e:
                logger.warning(f"    Could not restore data to message_attachements: {e}")

    logger.info("[Downgrade Step 2] Complete")
    logger.info("[Note] Kept derived_documents (old base schema), message_attachments/message_attachements")
