"""Add messages table

Revision ID: 40d4252deb5b
Revises: dc285eabea90
Create Date: 2025-07-31 11:25:30.567076

"""

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 = "40d4252deb5b"
down_revision: Union[str, None] = "dc285eabea90"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "messages",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("agent_instance_id", sa.UUID(), nullable=False),
        sa.Column(
            "sender_type", sa.Enum("AGENT", "USER", name="sendertype"), nullable=False
        ),
        sa.Column("content", sa.Text(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("requires_user_input", sa.Boolean(), nullable=False),
        sa.Column(
            "message_metadata", postgresql.JSONB(astext_type=sa.Text()), nullable=True
        ),
        sa.ForeignKeyConstraint(
            ["agent_instance_id"],
            ["agent_instances.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        "idx_messages_instance_created",
        "messages",
        ["agent_instance_id", "created_at"],
        unique=False,
    )

    # Migrate data from old tables to messages table
    # This combines AgentStep, AgentQuestion, and AgentUserFeedback into Messages
    # ordered by timestamp
    op.execute("""
        -- Insert agent steps as agent messages
        INSERT INTO messages (id, agent_instance_id, sender_type, content, created_at, requires_user_input, message_metadata)
        SELECT
            id,
            agent_instance_id,
            'AGENT',
            description,
            created_at,
            false,
            jsonb_build_object('source', 'agent_step', 'step_number', step_number)
        FROM agent_steps;

        -- Insert agent questions as agent messages
        INSERT INTO messages (id, agent_instance_id, sender_type, content, created_at, requires_user_input, message_metadata)
        SELECT
            id,
            agent_instance_id,
            'AGENT',
            question_text,
            asked_at,
            CASE WHEN answer_text IS NULL THEN true ELSE false END,
            jsonb_build_object('source', 'agent_question')
        FROM agent_questions;

        -- Insert question answers as user messages (only for answered questions)
        INSERT INTO messages (id, agent_instance_id, sender_type, content, created_at, requires_user_input, message_metadata)
        SELECT
            gen_random_uuid(),
            agent_instance_id,
            'USER',
            answer_text,
            answered_at,
            false,
            jsonb_build_object('source', 'question_answer', 'question_id', id::text, 'answered_by_user_id', answered_by_user_id::text)
        FROM agent_questions
        WHERE answer_text IS NOT NULL AND answered_at IS NOT NULL;

        -- Insert user feedback as user messages
        INSERT INTO messages (id, agent_instance_id, sender_type, content, created_at, requires_user_input, message_metadata)
        SELECT
            id,
            agent_instance_id,
            'USER',
            feedback_text,
            created_at,
            false,
            jsonb_build_object('source', 'user_feedback', 'created_by_user_id', created_by_user_id::text)
        FROM agent_user_feedback;

        -- Mark all agent instances as completed
        UPDATE agent_instances SET status = 'COMPLETED' WHERE status != 'COMPLETED';
    """)

    # Add last_read_message_id to agent_instances
    op.add_column(
        "agent_instances",
        sa.Column("last_read_message_id", postgresql.UUID(as_uuid=True), nullable=True),
    )
    # Create the foreign key constraint after both tables exist to avoid circular dependency
    op.create_foreign_key(
        "fk_agent_instances_last_read_message",
        "agent_instances",
        "messages",
        ["last_read_message_id"],
        ["id"],
        ondelete="SET NULL",
    )

    # Create combined function for message notifications
    op.execute("""
        CREATE OR REPLACE FUNCTION notify_message_change() RETURNS trigger AS $$
        DECLARE
            channel_name text;
            payload text;
            event_type text;
        BEGIN
            -- Create channel name based on instance ID
            channel_name := 'message_channel_' || NEW.agent_instance_id::text;

            -- Determine event type
            IF TG_OP = 'INSERT' THEN
                event_type := 'message_insert';
            ELSIF TG_OP = 'UPDATE' THEN
                event_type := 'message_update';
            END IF;

            -- Create JSON payload with message data
            payload := json_build_object(
                'event_type', event_type,
                'id', NEW.id,
                'agent_instance_id', NEW.agent_instance_id,
                'sender_type', NEW.sender_type,
                'content', NEW.content,
                'created_at', to_char(NEW.created_at AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'),
                'requires_user_input', NEW.requires_user_input,
                'message_metadata', NEW.message_metadata,
                'old_requires_user_input', CASE
                    WHEN TG_OP = 'UPDATE' THEN OLD.requires_user_input
                    ELSE NULL
                END
            )::text;

            -- Send notification (quote channel name for UUIDs with hyphens)
            EXECUTE format('NOTIFY %I, %L', channel_name, payload);

            RETURN NEW;
        END;
        $$ LANGUAGE plpgsql;
    """)

    # Create single trigger for both INSERT and UPDATE on messages table
    op.execute("""
        CREATE TRIGGER message_change_notify
        AFTER INSERT OR UPDATE ON messages
        FOR EACH ROW
        EXECUTE FUNCTION notify_message_change();
    """)

    # Create function for status change notifications
    op.execute("""
        CREATE OR REPLACE FUNCTION notify_status_change() RETURNS trigger AS $$
        DECLARE
            channel_name text;
            payload text;
        BEGIN
            -- Only notify if status actually changed
            IF OLD.status IS DISTINCT FROM NEW.status THEN
                -- Create channel name based on instance ID
                channel_name := 'message_channel_' || NEW.id::text;

                -- Create JSON payload with status update data
                payload := json_build_object(
                    'event_type', 'status_update',
                    'instance_id', NEW.id,
                    'status', NEW.status,
                    'timestamp', to_char(NOW() AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"')
                )::text;

                -- Send notification (quote channel name for UUIDs with hyphens)
                EXECUTE format('NOTIFY %I, %L', channel_name, payload);
            END IF;

            RETURN NEW;
        END;
        $$ LANGUAGE plpgsql;
    """)

    # Create trigger on agent_instances table for status updates
    op.execute("""
        CREATE TRIGGER agent_instance_status_notify
        AFTER UPDATE OF status ON agent_instances
        FOR EACH ROW
        EXECUTE FUNCTION notify_status_change();
    """)

    # Drop the old tables after migration
    op.drop_table("agent_user_feedback")
    op.drop_table("agent_questions")
    op.drop_table("agent_steps")
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    # Recreate the old tables first
    op.create_table(
        "agent_steps",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("agent_instance_id", sa.UUID(), nullable=False),
        sa.Column("step_number", sa.Integer(), nullable=False),
        sa.Column("description", sa.Text(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["agent_instance_id"],
            ["agent_instances.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )

    op.create_table(
        "agent_questions",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("agent_instance_id", sa.UUID(), nullable=False),
        sa.Column("question_text", sa.Text(), nullable=False),
        sa.Column("answer_text", sa.Text(), nullable=True),
        sa.Column("answered_by_user_id", sa.UUID(), nullable=True),
        sa.Column("asked_at", sa.DateTime(), nullable=False),
        sa.Column("answered_at", sa.DateTime(), nullable=True),
        sa.Column("is_active", sa.Boolean(), nullable=False),
        sa.ForeignKeyConstraint(
            ["agent_instance_id"],
            ["agent_instances.id"],
        ),
        sa.ForeignKeyConstraint(
            ["answered_by_user_id"],
            ["users.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )

    op.create_table(
        "agent_user_feedback",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("agent_instance_id", sa.UUID(), nullable=False),
        sa.Column("created_by_user_id", sa.UUID(), nullable=False),
        sa.Column("feedback_text", sa.Text(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("retrieved_at", sa.DateTime(), nullable=True),
        sa.ForeignKeyConstraint(
            ["agent_instance_id"],
            ["agent_instances.id"],
        ),
        sa.ForeignKeyConstraint(
            ["created_by_user_id"],
            ["users.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )

    # Note: We cannot fully restore the original data as some information is lost
    # (e.g., step_number ordering, which user created feedback vs answered questions)
    # This is a best-effort restoration

    # Drop status change trigger and function
    op.execute(
        "DROP TRIGGER IF EXISTS agent_instance_status_notify ON agent_instances;"
    )
    op.execute("DROP FUNCTION IF EXISTS notify_status_change();")

    # Drop message trigger and function
    op.execute("DROP TRIGGER IF EXISTS message_change_notify ON messages;")
    op.execute("DROP FUNCTION IF EXISTS notify_message_change();")

    op.drop_constraint(
        "fk_agent_instances_last_read_message", "agent_instances", type_="foreignkey"
    )
    op.drop_column("agent_instances", "last_read_message_id")
    op.drop_index("idx_messages_instance_created", table_name="messages")
    op.drop_table("messages")
    # ### end Alembic commands ###
