"""simplify_agent_architecture

Revision ID: 435419f8121c
Revises: 50ae013ce530
Create Date: 2026-03-04 11:01:11.127364

"""

from typing import Sequence, Union

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


# revision identifiers, used by Alembic.
revision: str = "435419f8121c"
down_revision: Union[str, Sequence[str], None] = "50ae013ce530"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # 1. Delete memories.agent_id dependencies
    op.drop_constraint("fk_memories_agent_id", "memories", type_="foreignkey")
    op.drop_constraint("chk_memory_type_agent_id", "memories", type_="check")
    op.execute("DROP INDEX IF EXISTS ix_memories_agent_type_status")
    op.drop_column("memories", "agent_id")

    # 2. Delete user_agents table
    _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_constraint("chk_agent_type", "user_agents", type_="check")
    op.drop_constraint(
        "uq_user_agents_user_id_agent_type", "user_agents", type_="unique"
    )

    op.execute("DROP INDEX IF EXISTS ix_user_agents_status")
    op.execute("DROP INDEX IF EXISTS ix_user_agents_agent_type")

    op.drop_table("user_agents")

    # 3. Rename user_agent_catalog to system_agents
    _drop_rls("user_agent_catalog")

    op.rename_table("user_agent_catalog", "system_agents")

    op.execute(
        "ALTER TABLE system_agents RENAME CONSTRAINT fk_user_agent_catalog_llm_id "
        "TO fk_system_agents_llm_id"
    )
    op.execute(
        "ALTER TABLE system_agents RENAME CONSTRAINT chk_user_agent_catalog_status "
        "TO chk_system_agents_status"
    )

    _enable_rls("system_agents")

    # 4. Update trigger
    op.execute("DROP TRIGGER IF EXISTS on_auth_user_created ON auth.users")
    op.execute("DROP FUNCTION IF EXISTS public.create_profile_for_new_user()")

    op.execute("""
        CREATE OR REPLACE FUNCTION public.create_profile_for_new_user()
        RETURNS trigger
        LANGUAGE plpgsql
        SECURITY DEFINER
        SET search_path = ''
        AS $$
        DECLARE
            invite_code_value TEXT;
            referrer_id UUID;
            new_code TEXT;
            attempts INT := 0;
        BEGIN
            INSERT INTO public.profiles (id, username, avatar_url, bio, settings, referred_by, 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,
                NULL,
                now(),
                now()
            )
            ON CONFLICT (id) DO NOTHING;

            LOOP
                BEGIN
                    new_code := public.generate_invite_code();
                    INSERT INTO public.invite_codes (code, owner_id, status, used_count, max_uses, expires_at, reward_config)
                    VALUES (
                        new_code,
                        NEW.id,
                        'active',
                        0,
                        NULL,
                        NULL,
                        '{}'::jsonb
                    );
                    EXIT;
                EXCEPTION WHEN unique_violation THEN
                    attempts := attempts + 1;
                    IF attempts >= 100 THEN
                        RAISE EXCEPTION 'Failed to generate unique invite code after 100 attempts';
                    END IF;
                END;
            END LOOP;

            invite_code_value := NEW.raw_user_meta_data ->> 'invite_code';
            IF invite_code_value IS NOT NULL AND length(invite_code_value) = 4 THEN
                invite_code_value := upper(invite_code_value);
                IF invite_code_value ~ '^[ABCDEFGHJKMNPQRSTUVWXYZ23456789]{4}$' THEN
                    UPDATE public.invite_codes
                    SET used_count = used_count + 1
                    WHERE code = invite_code_value
                        AND status = 'active'
                        AND (max_uses IS NULL OR used_count < max_uses)
                        AND (expires_at IS NULL OR expires_at > NOW())
                    RETURNING owner_id INTO referrer_id;

                    IF referrer_id IS NOT NULL THEN
                        UPDATE public.profiles
                        SET referred_by = referrer_id
                        WHERE id = NEW.id;
                    END IF;
                END IF;
            END IF;
            
            RETURN NEW;
        END;
        $$
    """)

    op.execute("""
        CREATE TRIGGER on_auth_user_created
            AFTER INSERT ON auth.users
            FOR EACH ROW EXECUTE FUNCTION public.create_profile_for_new_user()
    """)

    # 5. Update existing profiles.settings
    op.execute("""
        UPDATE profiles 
        SET settings = jsonb_set(
            COALESCE(settings, '{}'::jsonb),
            '{agent_prompts}',
            '{}'::jsonb
        )
        WHERE NOT settings ? 'agent_prompts'
    """)


def downgrade() -> None:
    # 1. Revert trigger
    op.execute("DROP TRIGGER IF EXISTS on_auth_user_created ON auth.users")
    op.execute("DROP FUNCTION IF EXISTS public.create_profile_for_new_user()")

    op.execute("""
        CREATE OR REPLACE FUNCTION public.create_profile_for_new_user()
        RETURNS trigger
        LANGUAGE plpgsql
        SECURITY DEFINER
        SET search_path = ''
        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;
        $$
    """)

    op.execute("""
        CREATE TRIGGER on_auth_user_created
            AFTER INSERT ON auth.users
            FOR EACH ROW EXECUTE FUNCTION public.create_profile_for_new_user()
    """)

    # 2. Revert rename: system_agents -> user_agent_catalog
    _drop_rls("system_agents")

    op.rename_table("system_agents", "user_agent_catalog")

    op.execute(
        "ALTER TABLE user_agent_catalog RENAME CONSTRAINT fk_system_agents_llm_id "
        "TO fk_user_agent_catalog_llm_id"
    )
    op.execute(
        "ALTER TABLE user_agent_catalog RENAME CONSTRAINT chk_system_agents_status "
        "TO chk_user_agent_catalog_status"
    )

    _enable_rls("user_agent_catalog")

    # 3. Recreate user_agents table
    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"),
    )

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

    op.execute("CREATE INDEX ix_user_agents_agent_type ON user_agents (agent_type)")
    op.execute("CREATE INDEX ix_user_agents_status ON user_agents (status)")

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

    # 4. Recreate memories.agent_id
    op.add_column("memories", sa.Column("agent_id", sa.UUID(), nullable=True))

    op.create_foreign_key(
        "fk_memories_agent_id",
        "memories",
        "user_agents",
        ["agent_id"],
        ["id"],
        ondelete="CASCADE",
    )

    op.execute(
        "CREATE INDEX ix_memories_agent_type_status ON memories (agent_id, memory_type, status)"
    )

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


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} "
            f"FOR SELECT TO {role} USING (false)"
        )
        op.execute(
            f"CREATE POLICY {role}_insert_{table_name} ON {table_name} "
            f"FOR INSERT TO {role} WITH CHECK (false)"
        )
        op.execute(
            f"CREATE POLICY {role}_update_{table_name} ON {table_name} "
            f"FOR UPDATE TO {role} USING (false) WITH CHECK (false)"
        )
        op.execute(
            f"CREATE POLICY {role}_delete_{table_name} ON {table_name} "
            f"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")