"""add_user_agent_catalog

Revision ID: 50ae013ce530
Revises: 202602270006
Create Date: 2026-03-02 15:34:25.995336

"""

from typing import Sequence, Union

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


revision: str = "50ae013ce530"
down_revision: Union[str, Sequence[str], None] = "202602270006"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    op.drop_constraint("uq_user_agents_user_id", "user_agents", type_="unique")

    op.create_unique_constraint(
        "uq_user_agents_user_id_agent_type", "user_agents", ["user_id", "agent_type"]
    )

    op.create_table(
        "user_agent_catalog",
        sa.Column("agent_type", sa.String(20), nullable=False),
        sa.Column("llm_id", sa.UUID(), nullable=False),
        sa.Column("status", sa.String(20), nullable=False),
        sa.Column(
            "config",
            postgresql.JSONB(astext_type=sa.Text()),
            server_default="{}",
            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("agent_type"),
        sa.ForeignKeyConstraint(
            ["llm_id"],
            ["llms.id"],
            name="fk_user_agent_catalog_llm_id",
            ondelete="RESTRICT",
        ),
    )

    op.execute(
        "ALTER TABLE user_agent_catalog "
        "ADD CONSTRAINT chk_user_agent_catalog_status "
        "CHECK (status IN ('active', 'paused', 'migrating'))"
    )

    _enable_rls("user_agent_catalog")

    op.execute("""
        CREATE OR REPLACE FUNCTION public.create_profile_for_new_user()
        RETURNS trigger
        LANGUAGE plpgsql
        SECURITY DEFINER
        SET search_path = public
        AS $$
        BEGIN
            INSERT INTO public.profiles (id, username, avatar_url, bio, settings, created_at, updated_at)
            VALUES (
                NEW.id,
                COALESCE(
                    NEW.raw_user_meta_data ->> 'username',
                    split_part(NEW.email, '@', 1),
                    'user_' || substring(NEW.id::text, 1, 8)
                ),
                NULL,
                NULL,
                '{}'::jsonb,
                now(),
                now()
            )
            ON CONFLICT (id) DO NOTHING;

            INSERT INTO public.user_agents (id, user_id, llm_id, agent_type, config, status, created_by, updated_by)
            SELECT 
                gen_random_uuid(),
                NEW.id,
                uac.llm_id,
                uac.agent_type,
                uac.config,
                uac.status,
                NEW.id,
                NEW.id
            FROM public.user_agent_catalog uac;
            
            RETURN NEW;
        END;
        $$;
    """)


def downgrade() -> None:
    op.execute("""
        CREATE OR REPLACE FUNCTION public.create_profile_for_new_user()
        RETURNS trigger
        LANGUAGE plpgsql
        SECURITY DEFINER
        SET search_path = public
        AS $$
        BEGIN
            INSERT INTO public.profiles (id, username, avatar_url, bio, settings, created_at, updated_at)
            VALUES (
                NEW.id,
                COALESCE(
                    NEW.raw_user_meta_data ->> 'username',
                    split_part(NEW.email, '@', 1),
                    'user_' || substring(NEW.id::text, 1, 8)
                ),
                NULL,
                NULL,
                '{}'::jsonb,
                now(),
                now()
            )
            ON CONFLICT (id) DO NOTHING;

            RETURN NEW;
        END;
        $$;
    """)

    op.drop_constraint(
        "uq_user_agents_user_id_agent_type", "user_agents", type_="unique"
    )

    op.create_unique_constraint("uq_user_agents_user_id", "user_agents", ["user_id"])

    _drop_rls("user_agent_catalog")
    op.drop_constraint(
        "chk_user_agent_catalog_status", "user_agent_catalog", type_="check"
    )
    op.drop_constraint(
        "fk_user_agent_catalog_llm_id", "user_agent_catalog", type_="foreignkey"
    )
    op.drop_table("user_agent_catalog")


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")
