"""add exploration_sessions and exploration_messages tables

Revision ID: 0e4a7c91d2b3
Revises: f2a9b7c34d18
Create Date: 2026-08-03
"""

from typing import Sequence, Union

import sqlalchemy as sa
import sqlmodel  # noqa: F401
from alembic import op

revision: str = "0e4a7c91d2b3"
down_revision: Union[str, Sequence[str], None] = "f2a9b7c34d18"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    op.create_table(
        "exploration_sessions",
        sa.Column("id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("campaign_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("target_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("persona", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("goal", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("seed_ref", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("status", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("triggered_by", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("experience", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("judge_review", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("turn_count", sa.Integer(), nullable=False),
        sa.Column("error", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=True),
        sa.Column("closed_at", sa.DateTime(), nullable=True),
        sa.PrimaryKeyConstraint("id"),
        sa.ForeignKeyConstraint(["campaign_id"], ["campaigns.id"]),
        sa.ForeignKeyConstraint(["target_id"], ["eval_targets.id"]),
    )
    op.create_index(
        "ix_exploration_sessions_campaign_id",
        "exploration_sessions",
        ["campaign_id"],
    )
    op.create_table(
        "exploration_messages",
        sa.Column("id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("session_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("round_index", sa.Integer(), nullable=False),
        sa.Column("role", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("content", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("latency_ms", sa.Integer(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=True),
        sa.PrimaryKeyConstraint("id"),
        sa.ForeignKeyConstraint(["session_id"], ["exploration_sessions.id"]),
    )
    op.create_index(
        "ix_exploration_messages_session_id",
        "exploration_messages",
        ["session_id"],
    )


def downgrade() -> None:
    op.drop_index("ix_exploration_messages_session_id")
    op.drop_table("exploration_messages")
    op.drop_index("ix_exploration_sessions_campaign_id")
    op.drop_table("exploration_sessions")
