"""Persist chat state, indexing tasks and fixed batch membership.

Revision ID: 0003
Revises: 0002
"""
from alembic import op
import sqlalchemy as sa

revision = "0003"
down_revision = "0002"
branch_labels = None
depends_on = None


def upgrade():
    inspector = sa.inspect(op.get_bind())
    tables = set(inspector.get_table_names())
    if "project_ids" not in {column["name"] for column in inspector.get_columns("scoring_batches")}:
        op.add_column("scoring_batches", sa.Column("project_ids", sa.JSON(), nullable=True))
    if "index_tasks" not in tables:
        op.create_table("index_tasks",
        sa.Column("id", sa.Integer(), primary_key=True),
        sa.Column("course_version_id", sa.Integer(), sa.ForeignKey("course_versions.id"), nullable=False),
        sa.Column("submission_file_id", sa.Integer(), sa.ForeignKey("submission_files.id")),
        sa.Column("kind", sa.String(30), nullable=False), sa.Column("status", sa.String(30), nullable=False),
        sa.Column("result", sa.JSON(), nullable=False), sa.Column("error", sa.Text()),
        sa.Column("created_at", sa.DateTime(timezone=True)), sa.Column("updated_at", sa.DateTime(timezone=True)))
    if "conversations" not in tables:
        op.create_table("conversations",
        sa.Column("id", sa.Integer(), primary_key=True), sa.Column("title", sa.String(200), nullable=False),
        sa.Column("course_version_id", sa.Integer(), sa.ForeignKey("course_versions.id")),
        sa.Column("batch_id", sa.Integer(), sa.ForeignKey("scoring_batches.id")),
        sa.Column("project_id", sa.Integer(), sa.ForeignKey("projects.id")),
        sa.Column("created_at", sa.DateTime(timezone=True)), sa.Column("updated_at", sa.DateTime(timezone=True)))
    if "chat_messages" not in tables:
        op.create_table("chat_messages",
        sa.Column("id", sa.Integer(), primary_key=True),
        sa.Column("conversation_id", sa.Integer(), sa.ForeignKey("conversations.id"), nullable=False),
        sa.Column("role", sa.String(20), nullable=False), sa.Column("content", sa.Text(), nullable=False),
        sa.Column("citations", sa.JSON(), nullable=False), sa.Column("tool_calls", sa.JSON(), nullable=False),
        sa.Column("created_at", sa.DateTime(timezone=True)))
    if "agent_runs" not in tables:
        op.create_table("agent_runs", sa.Column("id", sa.Integer(), primary_key=True),
        sa.Column("conversation_id", sa.Integer(), sa.ForeignKey("conversations.id"), nullable=False),
        sa.Column("status", sa.String(30), nullable=False), sa.Column("error", sa.Text()),
        sa.Column("created_at", sa.DateTime(timezone=True)), sa.Column("finished_at", sa.DateTime(timezone=True)))
    if "pending_actions" not in tables:
        op.create_table("pending_actions", sa.Column("id", sa.Integer(), primary_key=True),
        sa.Column("conversation_id", sa.Integer(), sa.ForeignKey("conversations.id"), nullable=False),
        sa.Column("kind", sa.String(40), nullable=False), sa.Column("payload", sa.JSON(), nullable=False),
        sa.Column("preview", sa.JSON(), nullable=False), sa.Column("expected", sa.JSON(), nullable=False),
        sa.Column("idempotency_key", sa.String(64), nullable=False, unique=True),
        sa.Column("status", sa.String(30), nullable=False), sa.Column("result", sa.JSON()),
        sa.Column("created_at", sa.DateTime(timezone=True)), sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False))
    if "ix_chat_messages_conversation_id" not in {index["name"] for index in sa.inspect(op.get_bind()).get_indexes("chat_messages")}:
        op.create_index("ix_chat_messages_conversation_id", "chat_messages", ["conversation_id"])
    if "ix_pending_actions_conversation_id" not in {index["name"] for index in sa.inspect(op.get_bind()).get_indexes("pending_actions")}:
        op.create_index("ix_pending_actions_conversation_id", "pending_actions", ["conversation_id"])


def downgrade():
    op.drop_index("ix_pending_actions_conversation_id", table_name="pending_actions")
    op.drop_index("ix_chat_messages_conversation_id", table_name="chat_messages")
    for table in ("pending_actions", "agent_runs", "chat_messages", "conversations", "index_tasks"):
        op.drop_table(table)
    op.drop_column("scoring_batches", "project_ids")