"""invite_codes table and profile referral fields

Revision ID: 202602270006
Revises: 202602260005
Create Date: 2026-02-27 10:00:00
"""

from typing import Sequence, Union

from alembic import op

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


def upgrade() -> None:
    op.execute(
        """
        CREATE TABLE invite_codes (
            id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            code VARCHAR(8) NOT NULL UNIQUE CHECK (code ~ '^[ABCDEFGHJKMNPQRSTUVWXYZ23456789]{8}$'),
            owner_id UUID REFERENCES profiles(id) ON DELETE SET NULL,
            status VARCHAR(20) NOT NULL DEFAULT 'active' CHECK (status IN ('active', 'disabled', 'expired')),
            used_count INTEGER NOT NULL DEFAULT 0 CHECK (used_count >= 0),
            max_uses INTEGER CHECK (max_uses IS NULL OR max_uses >= 1),
            expires_at TIMESTAMPTZ NULL,
            reward_config JSONB NOT NULL DEFAULT '{}',
            created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
        )
        """
    )
    op.execute("CREATE INDEX ix_invite_codes_owner_id ON invite_codes(owner_id)")
    op.execute(
        "CREATE INDEX ix_invite_codes_code ON invite_codes(code) WHERE status = 'active'"
    )

    op.execute("ALTER TABLE invite_codes ENABLE ROW LEVEL SECURITY")
    op.execute("DROP POLICY IF EXISTS invite_codes_all_denied ON invite_codes")
    op.execute(
        "CREATE POLICY invite_codes_all_denied ON invite_codes FOR ALL USING (false)"
    )

    op.execute(
        """
        ALTER TABLE profiles ADD COLUMN referred_by UUID REFERENCES profiles(id) ON DELETE SET NULL
        """
    )
    op.execute("CREATE INDEX ix_profiles_referred_by ON profiles(referred_by)")

    op.execute(
        """
        CREATE OR REPLACE FUNCTION public.generate_invite_code()
        RETURNS TEXT
        LANGUAGE plpgsql
        SECURITY DEFINER
        SET search_path = public
        AS $$
        DECLARE
            chars TEXT := 'ABCDEFGHJKMNPQRSTUVWXYZ23456789';
            result TEXT := '';
            i INT;
        BEGIN
            FOR i IN 1..8 LOOP
                result := result || substr(chars, floor(random() * length(chars) + 1)::int, 1);
            END LOOP;
            RETURN result;
        END;
        $$;
        """
    )

    op.execute(
        """
        CREATE OR REPLACE FUNCTION public.create_profile_for_new_user()
        RETURNS trigger
        LANGUAGE plpgsql
        SECURITY DEFINER
        SET search_path = public
        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) = 8 THEN
                invite_code_value := upper(invite_code_value);
                IF invite_code_value ~ '^[ABCDEFGHJKMNPQRSTUVWXYZ23456789]{8}$' 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;
        $$;
        """
    )


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.execute("DROP FUNCTION IF EXISTS public.generate_invite_code()")

    op.execute("DROP INDEX IF EXISTS ix_profiles_referred_by")
    op.execute("ALTER TABLE profiles DROP COLUMN IF EXISTS referred_by")

    op.execute("DROP POLICY IF EXISTS invite_codes_all_denied ON invite_codes")
    op.execute("ALTER TABLE invite_codes DISABLE ROW LEVEL SECURITY")
    op.execute("DROP INDEX IF EXISTS ix_invite_codes_code")
    op.execute("DROP INDEX IF EXISTS ix_invite_codes_owner_id")
    op.execute("DROP TABLE IF EXISTS invite_codes")
