"""add mounted ai runtime tables

Revision ID: 20260405_add_mounted_ai_runtime_tables
Revises: 20260331_add_qa_record_runtime_fields
Create Date: 2026-04-05
"""

from __future__ import annotations

from alembic import op
import sqlalchemy as sa

try:
    from sqlalchemy.dialects.postgresql import JSONB
except Exception:  # pragma: no cover
    JSONB = sa.JSON


revision = "20260405_add_mounted_ai_runtime_tables"
down_revision = "20260331_add_qa_record_runtime_fields"
branch_labels = None
depends_on = None


def _json_type(bind):
    return JSONB() if bind.dialect.name == "postgresql" else sa.JSON()


def _table_exists(bind, table_name: str) -> bool:
    return sa.inspect(bind).has_table(table_name)


def _index_exists(bind, table_name: str, index_name: str) -> bool:
    inspector = sa.inspect(bind)
    try:
        indexes = inspector.get_indexes(table_name)
    except Exception:
        return False
    return any(str(index.get("name")) == index_name for index in indexes)


def _create_index_if_missing(bind, table_name: str, index_name: str, columns: list[str], *, unique: bool = False) -> None:
    if not _index_exists(bind, table_name, index_name):
        op.create_index(index_name, table_name, columns, unique=unique)


def _base_columns() -> list[sa.Column]:
    return [
        sa.Column("tenant_id", sa.BigInteger(), nullable=False),
        sa.Column("created_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now()),
        sa.Column("updated_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now()),
    ]


def _create_mounted_topic_workspaces(bind) -> None:
    if _table_exists(bind, "mounted_topic_workspaces"):
        return
    op.create_table(
        "mounted_topic_workspaces",
        sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
        sa.Column("topic_key", sa.VARCHAR(length=128), nullable=False),
        sa.Column("source_system", sa.VARCHAR(length=32), nullable=False),
        sa.Column("channel_ref", sa.VARCHAR(length=128), nullable=True),
        sa.Column("thread_key", sa.VARCHAR(length=255), nullable=True),
        sa.Column("root_message_id", sa.VARCHAR(length=128), nullable=True),
        sa.Column("mode", sa.VARCHAR(length=24), nullable=False, server_default="chat"),
        sa.Column("status", sa.VARCHAR(length=24), nullable=False, server_default="active"),
        sa.Column("facts_json", _json_type(bind), nullable=True),
        sa.Column("goals_json", _json_type(bind), nullable=True),
        sa.Column("constraints_json", _json_type(bind), nullable=True),
        sa.Column("options_json", _json_type(bind), nullable=True),
        sa.Column("open_questions", _json_type(bind), nullable=True),
        sa.Column("pending_actions", _json_type(bind), nullable=True),
        sa.Column("bound_entities", _json_type(bind), nullable=True),
        sa.Column("readiness_score", sa.Float(), nullable=True, server_default="0"),
        sa.Column("last_message_id", sa.VARCHAR(length=128), nullable=True),
        sa.Column("last_triggered_at", sa.TIMESTAMP(), nullable=True),
        sa.Column("metadata_json", _json_type(bind), nullable=True),
        *_base_columns(),
        sa.UniqueConstraint(
            "tenant_id",
            "source_system",
            "topic_key",
            name="uq_mounted_topic_workspace_source_topic",
        ),
    )
    for index_name, columns in (
        ("ix_mounted_topic_workspaces_tenant_id", ["tenant_id"]),
        ("ix_mounted_topic_workspaces_topic_key", ["topic_key"]),
        ("ix_mounted_topic_workspaces_source_system", ["source_system"]),
        ("ix_mounted_topic_workspaces_channel_ref", ["channel_ref"]),
        ("ix_mounted_topic_workspaces_thread_key", ["thread_key"]),
        ("ix_mounted_topic_workspaces_root_message_id", ["root_message_id"]),
        ("ix_mounted_topic_workspaces_mode", ["mode"]),
        ("ix_mounted_topic_workspaces_status", ["status"]),
        ("ix_mounted_topic_workspaces_last_message_id", ["last_message_id"]),
    ):
        _create_index_if_missing(bind, "mounted_topic_workspaces", index_name, columns)


def _create_mounted_draft_workspaces(bind) -> None:
    if _table_exists(bind, "mounted_draft_workspaces"):
        return
    op.create_table(
        "mounted_draft_workspaces",
        sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
        sa.Column(
            "topic_workspace_id",
            sa.Integer(),
            sa.ForeignKey("mounted_topic_workspaces.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column("state", sa.VARCHAR(length=24), nullable=False, server_default="collecting"),
        sa.Column("draft_version", sa.Integer(), nullable=False, server_default="0"),
        sa.Column("requirement_card", _json_type(bind), nullable=True),
        sa.Column("draft_payload", _json_type(bind), nullable=True),
        sa.Column("required_confirmations", _json_type(bind), nullable=True),
        sa.Column("last_ready_at", sa.TIMESTAMP(), nullable=True),
        sa.Column("confirm_message_id", sa.VARCHAR(length=128), nullable=True),
        *_base_columns(),
        sa.UniqueConstraint(
            "tenant_id",
            "topic_workspace_id",
            name="uq_mounted_draft_workspace_topic",
        ),
    )
    for index_name, columns in (
        ("ix_mounted_draft_workspaces_tenant_id", ["tenant_id"]),
        ("ix_mounted_draft_workspaces_topic_workspace_id", ["topic_workspace_id"]),
        ("ix_mounted_draft_workspaces_state", ["state"]),
        ("ix_mounted_draft_workspaces_confirm_message_id", ["confirm_message_id"]),
    ):
        _create_index_if_missing(bind, "mounted_draft_workspaces", index_name, columns)


def _create_mounted_conversation_runs(bind) -> None:
    if _table_exists(bind, "mounted_conversation_runs"):
        return
    op.create_table(
        "mounted_conversation_runs",
        sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
        sa.Column("trace_id", sa.VARCHAR(length=64), nullable=False),
        sa.Column(
            "topic_workspace_id",
            sa.Integer(),
            sa.ForeignKey("mounted_topic_workspaces.id", ondelete="SET NULL"),
            nullable=True,
        ),
        sa.Column("source_system", sa.VARCHAR(length=32), nullable=False),
        sa.Column("source_message_id", sa.VARCHAR(length=128), nullable=True),
        sa.Column("mode_before", sa.VARCHAR(length=24), nullable=True),
        sa.Column("mode_after", sa.VARCHAR(length=24), nullable=True),
        sa.Column("active_agent", sa.VARCHAR(length=64), nullable=True),
        sa.Column("status", sa.VARCHAR(length=24), nullable=False, server_default="queued"),
        sa.Column("trigger_reason", sa.VARCHAR(length=128), nullable=True),
        sa.Column("context_layers_hit", _json_type(bind), nullable=True),
        sa.Column("tool_calls", _json_type(bind), nullable=True),
        sa.Column("final_reply", sa.Text(), nullable=True),
        sa.Column("execution_outcome", _json_type(bind), nullable=True),
        sa.Column("error_text", sa.Text(), nullable=True),
        sa.Column("started_at", sa.TIMESTAMP(), nullable=True),
        sa.Column("completed_at", sa.TIMESTAMP(), nullable=True),
        *_base_columns(),
        sa.UniqueConstraint("tenant_id", "trace_id", name="uq_mounted_conversation_run_trace"),
    )
    for index_name, columns in (
        ("ix_mounted_conversation_runs_tenant_id", ["tenant_id"]),
        ("ix_mounted_conversation_runs_trace_id", ["trace_id"]),
        ("ix_mounted_conversation_runs_topic_workspace_id", ["topic_workspace_id"]),
        ("ix_mounted_conversation_runs_source_system", ["source_system"]),
        ("ix_mounted_conversation_runs_source_message_id", ["source_message_id"]),
        ("ix_mounted_conversation_runs_mode_before", ["mode_before"]),
        ("ix_mounted_conversation_runs_mode_after", ["mode_after"]),
        ("ix_mounted_conversation_runs_active_agent", ["active_agent"]),
        ("ix_mounted_conversation_runs_status", ["status"]),
        ("ix_mounted_conversation_runs_trigger_reason", ["trigger_reason"]),
    ):
        _create_index_if_missing(bind, "mounted_conversation_runs", index_name, columns)


def _create_mounted_agent_traces(bind) -> None:
    if _table_exists(bind, "mounted_agent_traces"):
        return
    op.create_table(
        "mounted_agent_traces",
        sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
        sa.Column(
            "conversation_run_id",
            sa.Integer(),
            sa.ForeignKey("mounted_conversation_runs.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column("agent_name", sa.VARCHAR(length=64), nullable=False),
        sa.Column("status", sa.VARCHAR(length=24), nullable=False, server_default="completed"),
        sa.Column("latency_ms", sa.Integer(), nullable=True),
        sa.Column("input_summary", _json_type(bind), nullable=True),
        sa.Column("output_summary", _json_type(bind), nullable=True),
        sa.Column("tool_trace", _json_type(bind), nullable=True),
        sa.Column("failure_reason", sa.Text(), nullable=True),
        *_base_columns(),
    )
    for index_name, columns in (
        ("ix_mounted_agent_traces_tenant_id", ["tenant_id"]),
        ("ix_mounted_agent_traces_conversation_run_id", ["conversation_run_id"]),
        ("ix_mounted_agent_traces_agent_name", ["agent_name"]),
        ("ix_mounted_agent_traces_status", ["status"]),
    ):
        _create_index_if_missing(bind, "mounted_agent_traces", index_name, columns)


def _create_mounted_tool_capability_policies(bind) -> None:
    if _table_exists(bind, "mounted_tool_capability_policies"):
        return
    op.create_table(
        "mounted_tool_capability_policies",
        sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
        sa.Column("tool_name", sa.VARCHAR(length=128), nullable=False),
        sa.Column("risk_level", sa.VARCHAR(length=24), nullable=False, server_default="medium"),
        sa.Column("confirm_required", sa.Boolean(), nullable=False, server_default=sa.true()),
        sa.Column("auto_execute_allowed", sa.Boolean(), nullable=False, server_default=sa.false()),
        sa.Column("is_active", sa.Boolean(), nullable=False, server_default=sa.true()),
        sa.Column("source_scope", _json_type(bind), nullable=True),
        sa.Column("tenant_scope", _json_type(bind), nullable=True),
        sa.Column("policy_metadata", _json_type(bind), nullable=True),
        sa.Column("notes", sa.Text(), nullable=True),
        *_base_columns(),
        sa.UniqueConstraint(
            "tenant_id",
            "tool_name",
            name="uq_mounted_tool_capability_policy_tool",
        ),
    )
    for index_name, columns in (
        ("ix_mounted_tool_capability_policies_tenant_id", ["tenant_id"]),
        ("ix_mounted_tool_capability_policies_tool_name", ["tool_name"]),
        ("ix_mounted_tool_capability_policies_risk_level", ["risk_level"]),
        ("ix_mounted_tool_capability_policies_is_active", ["is_active"]),
    ):
        _create_index_if_missing(bind, "mounted_tool_capability_policies", index_name, columns)


def upgrade() -> None:
    bind = op.get_bind()
    _create_mounted_topic_workspaces(bind)
    _create_mounted_draft_workspaces(bind)
    _create_mounted_conversation_runs(bind)
    _create_mounted_agent_traces(bind)
    _create_mounted_tool_capability_policies(bind)


def downgrade() -> None:
    bind = op.get_bind()
    for table_name in (
        "mounted_tool_capability_policies",
        "mounted_agent_traces",
        "mounted_conversation_runs",
        "mounted_draft_workspaces",
        "mounted_topic_workspaces",
    ):
        if _table_exists(bind, table_name):
            op.drop_table(table_name)