"""initial schema part 2: session and agent tables

Revision ID: 202602260002
Revises: 202602260001
Create Date: 2026-02-26 20:11:00
"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

revision: str = "202602260002"
down_revision: Union[str, Sequence[str], None] = "202602260001"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    op.create_table(
        "sessions",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("user_id", sa.UUID(), nullable=False),
        sa.Column("session_type", sa.String(length=20), nullable=False),
        sa.Column("job_id", sa.UUID(), nullable=True),
        sa.Column("title", sa.String(length=255), nullable=True),
        sa.Column("status", sa.String(length=20), nullable=False),
        sa.Column(
            "last_activity_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "message_count", sa.Integer(), server_default=sa.text("0"), nullable=False
        ),
        sa.Column(
            "total_tokens", sa.Integer(), server_default=sa.text("0"), nullable=False
        ),
        sa.Column(
            "total_cost",
            sa.Numeric(precision=12, scale=6),
            server_default=sa.text("0"),
            nullable=False,
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("ix_sessions_job_id", "sessions", ["job_id"], unique=False)
    op.execute(
        "CREATE INDEX ix_sessions_user_session_type_last_activity ON sessions (user_id, session_type, last_activity_at DESC)"
    )
    op.execute(
        "ALTER TABLE sessions ADD CONSTRAINT chk_session_type CHECK (session_type IN ('chat', 'automation'))"
    )
    op.execute(
        "ALTER TABLE sessions ADD CONSTRAINT chk_sessions_type_job_consistency CHECK ((session_type = 'chat' AND job_id IS NULL) OR (session_type = 'automation' AND job_id IS NOT NULL))"
    )
    op.execute(
        "ALTER TABLE sessions ADD CONSTRAINT chk_sessions_status CHECK (status IN ('pending', 'running', 'completed', 'failed'))"
    )
    op.create_foreign_key(
        "fk_sessions_user_id",
        "sessions",
        "users",
        ["user_id"],
        ["id"],
        referent_schema="auth",
        ondelete="CASCADE",
    )
    op.create_foreign_key(
        "fk_sessions_job_id_user_id",
        "sessions",
        "automation_jobs",
        ["job_id", "user_id"],
        ["id", "owner_id"],
        ondelete="RESTRICT",
    )
    op.create_index(
        "ix_sessions_job_user",
        "sessions",
        ["job_id", "user_id"],
        unique=False,
    )
    _enable_rls("sessions")

    op.create_table(
        "messages",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("session_id", sa.UUID(), nullable=False),
        sa.Column("seq", sa.Integer(), nullable=False),
        sa.Column("role", sa.String(length=20), nullable=False),
        sa.Column("content", sa.Text(), nullable=False),
        sa.Column("model_code", sa.String(length=50), nullable=True),
        sa.Column("tool_name", sa.String(length=100), nullable=True),
        sa.Column(
            "input_tokens", sa.Integer(), nullable=False, server_default=sa.text("0")
        ),
        sa.Column(
            "output_tokens", sa.Integer(), nullable=False, server_default=sa.text("0")
        ),
        sa.Column(
            "cost",
            sa.Numeric(precision=12, scale=6),
            nullable=False,
            server_default=sa.text("0"),
        ),
        sa.Column(
            "currency",
            sa.String(length=3),
            nullable=False,
            server_default=sa.text("'USD'"),
        ),
        sa.Column("latency_ms", sa.Integer(), nullable=True),
        sa.Column("metadata", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("session_id", "seq", name="uq_messages_session_seq"),
    )
    op.create_index("ix_messages_session_id", "messages", ["session_id"], unique=False)
    op.execute(
        "ALTER TABLE messages ADD CONSTRAINT chk_message_role CHECK (role IN ('user', 'assistant', 'system', 'tool'))"
    )
    op.create_foreign_key(
        "fk_messages_session_id",
        "messages",
        "sessions",
        ["session_id"],
        ["id"],
        ondelete="CASCADE",
    )
    _enable_rls("messages")

    op.create_table(
        "user_agents",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("user_id", sa.UUID(), nullable=False),
        sa.Column("llm_id", sa.UUID(), nullable=False),
        sa.Column("agent_type", sa.String(length=20), nullable=False),
        sa.Column(
            "config",
            postgresql.JSONB(astext_type=sa.Text()),
            server_default="{}",
            nullable=False,
        ),
        sa.Column("status", sa.String(length=20), nullable=False),
        sa.Column("created_by", sa.UUID(), nullable=True),
        sa.Column("updated_by", sa.UUID(), nullable=True),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("user_id", name="uq_user_agents_user_id"),
    )
    op.create_index(
        "ix_user_agents_agent_type", "user_agents", ["agent_type"], unique=False
    )
    op.create_index("ix_user_agents_status", "user_agents", ["status"], unique=False)
    op.execute(
        "ALTER TABLE user_agents ADD CONSTRAINT chk_agent_type CHECK (agent_type IN ('INTENT_RECOGNITION', 'TASK_EXECUTION', 'RESULT_REPORTING'))"
    )
    op.create_foreign_key(
        "fk_user_agents_user_id",
        "user_agents",
        "users",
        ["user_id"],
        ["id"],
        referent_schema="auth",
        ondelete="CASCADE",
    )
    op.create_foreign_key(
        "fk_user_agents_llm_id",
        "user_agents",
        "llms",
        ["llm_id"],
        ["id"],
        ondelete="RESTRICT",
    )
    op.create_foreign_key(
        "fk_user_agents_created_by",
        "user_agents",
        "users",
        ["created_by"],
        ["id"],
        referent_schema="auth",
        ondelete="SET NULL",
    )
    op.create_foreign_key(
        "fk_user_agents_updated_by",
        "user_agents",
        "users",
        ["updated_by"],
        ["id"],
        referent_schema="auth",
        ondelete="SET NULL",
    )
    _enable_rls("user_agents")

    op.create_table(
        "memories",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("owner_id", sa.UUID(), nullable=False),
        sa.Column("agent_id", sa.UUID(), nullable=True),
        sa.Column("memory_type", sa.String(length=20), nullable=False),
        sa.Column("title", sa.String(length=255), nullable=True),
        sa.Column("content", postgresql.JSONB(astext_type=sa.Text()), nullable=False),
        sa.Column("source", sa.String(length=20), nullable=False),
        sa.Column("status", sa.String(length=20), nullable=False),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        "ix_memories_owner_type_status",
        "memories",
        ["owner_id", "memory_type", "status"],
        unique=False,
    )
    op.create_index(
        "ix_memories_agent_type_status",
        "memories",
        ["agent_id", "memory_type", "status"],
        unique=False,
    )
    op.execute(
        "ALTER TABLE memories ADD CONSTRAINT chk_memory_type_agent_id CHECK ((memory_type = 'work' AND agent_id IS NOT NULL) OR (memory_type = 'user' AND agent_id IS NULL))"
    )
    op.create_foreign_key(
        "fk_memories_owner_id",
        "memories",
        "users",
        ["owner_id"],
        ["id"],
        referent_schema="auth",
        ondelete="CASCADE",
    )
    op.create_foreign_key(
        "fk_memories_agent_id",
        "memories",
        "user_agents",
        ["agent_id"],
        ["id"],
        ondelete="CASCADE",
    )
    _enable_rls("memories")


def downgrade() -> None:
    _drop_rls("memories")
    op.drop_constraint("fk_memories_agent_id", "memories", type_="foreignkey")
    op.drop_constraint("fk_memories_owner_id", "memories", type_="foreignkey")
    op.drop_table("memories")

    _drop_rls("user_agents")
    op.drop_constraint("fk_user_agents_updated_by", "user_agents", type_="foreignkey")
    op.drop_constraint("fk_user_agents_created_by", "user_agents", type_="foreignkey")
    op.drop_constraint("fk_user_agents_llm_id", "user_agents", type_="foreignkey")
    op.drop_constraint("fk_user_agents_user_id", "user_agents", type_="foreignkey")
    op.drop_table("user_agents")

    _drop_rls("messages")
    op.drop_constraint("fk_messages_session_id", "messages", type_="foreignkey")
    op.drop_table("messages")

    _drop_rls("sessions")
    op.drop_constraint("fk_sessions_job_id_user_id", "sessions", type_="foreignkey")
    op.drop_constraint("fk_sessions_user_id", "sessions", type_="foreignkey")
    op.drop_table("sessions")


def _enable_rls(table_name: str) -> None:
    for role in ["anon", "authenticated"]:
        for action in ["select", "insert", "update", "delete"]:
            op.execute(
                f"DROP POLICY IF EXISTS {role}_{action}_{table_name} ON {table_name}"
            )
    op.execute(f"ALTER TABLE {table_name} ENABLE ROW LEVEL SECURITY")
    for role in ["anon", "authenticated"]:
        op.execute(
            f"CREATE POLICY {role}_select_{table_name} ON {table_name} FOR SELECT TO {role} USING (false)"
        )
        op.execute(
            f"CREATE POLICY {role}_insert_{table_name} ON {table_name} FOR INSERT TO {role} WITH CHECK (false)"
        )
        op.execute(
            f"CREATE POLICY {role}_update_{table_name} ON {table_name} FOR UPDATE TO {role} USING (false) WITH CHECK (false)"
        )
        op.execute(
            f"CREATE POLICY {role}_delete_{table_name} ON {table_name} FOR DELETE TO {role} USING (false)"
        )


def _drop_rls(table_name: str) -> None:
    for role in ["anon", "authenticated"]:
        op.execute(f"DROP POLICY IF EXISTS {role}_delete_{table_name} ON {table_name}")
        op.execute(f"DROP POLICY IF EXISTS {role}_update_{table_name} ON {table_name}")
        op.execute(f"DROP POLICY IF EXISTS {role}_insert_{table_name} ON {table_name}")
        op.execute(f"DROP POLICY IF EXISTS {role}_select_{table_name} ON {table_name}")
    op.execute(f"ALTER TABLE {table_name} DISABLE ROW LEVEL SECURITY")
