"""add invite_codes table and profiles referred_by

Revision ID: 20260407_0002
Revises: 20260407_0001
Create Date: 2026-04-07 00:00:00
"""

from typing import Sequence, Union

from alembic import op

revision: str = "20260407_0002"
down_revision: Union[str, Sequence[str], None] = "20260407_0001"
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(6) NOT NULL UNIQUE CHECK (code ~ '^[ABCDEFGHJKMNPQRSTUVWXYZ23456789]{6}$'),
            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 = ''
        AS $$
        DECLARE
            chars TEXT := 'ABCDEFGHJKMNPQRSTUVWXYZ23456789';
            result TEXT := '';
            i INT;
        BEGIN
            FOR i IN 1..6 LOOP
                result := result || substr(chars, floor(random() * length(chars) + 1)::int, 1);
            END LOOP;
            RETURN result;
        END;
        $$;
        """
    )

    op.execute(
        """
        CREATE OR REPLACE FUNCTION public.initialize_profile_and_invite_code_on_signup()
        RETURNS trigger
        LANGUAGE plpgsql
        SECURITY DEFINER
        SET search_path = public
        AS $$
        DECLARE
            v_username text;
            v_ledger_id uuid;
            v_event_id text;
            v_invite_code text;
            v_referrer_id uuid;
            v_attempts int := 0;
            invite_code_value text;
        BEGIN
            v_username := 'user_' || substring(md5(new.id::text || clock_timestamp()::text || random()::text) from 1 for 6);
            v_ledger_id := md5(new.id::text || 'ledger' || clock_timestamp()::text || random()::text)::uuid;
            v_event_id := 'register:' || new.id::text;

            INSERT INTO public.profiles (id, username, avatar_url, bio, settings)
            VALUES (
                new.id,
                v_username,
                null,
                null,
                jsonb_build_object(
                    'version', 1,
                    'preferences', jsonb_build_object(
                        'interface_language', 'zh-CN',
                        'ai_language', 'zh-CN',
                        'timezone', 'Asia/Shanghai',
                        'country', 'CN'
                    ),
                    'privacy', jsonb_build_object('profile_visibility', 'public'),
                    'notification', jsonb_build_object(
                        'allow_notifications', true,
                        'allow_vibration', true
                    )
                )
            )
            ON CONFLICT (id) DO NOTHING;

            INSERT INTO public.user_points (
                user_id,
                balance,
                frozen_balance,
                lifetime_earned,
                lifetime_spent,
                version
            )
            VALUES (new.id, 100, 0, 100, 0, 0)
            ON CONFLICT (user_id) DO NOTHING;

            INSERT INTO public.points_ledger (
                id,
                user_id,
                direction,
                amount,
                balance_after,
                change_type,
                biz_type,
                biz_id,
                event_id,
                operator_id,
                metadata
            )
            VALUES (
                v_ledger_id,
                new.id,
                1,
                100,
                100,
                'register',
                null,
                null,
                v_event_id,
                null,
                jsonb_build_object(
                    'schema_version', 1,
                    'reason_code', 'REGISTER_WELCOME',
                    'operator_type', 'system',
                    'run_id', v_event_id,
                    'request_id', null,
                    'ext', jsonb_build_object('source', 'auth_signup')
                )
            )
            ON CONFLICT (user_id, event_id) DO NOTHING;

            LOOP
                BEGIN
                    v_invite_code := public.generate_invite_code();
                    INSERT INTO public.invite_codes (code, owner_id, status, used_count, max_uses, expires_at, reward_config)
                    VALUES (
                        v_invite_code,
                        new.id,
                        'active',
                        0,
                        NULL,
                        NULL,
                        '{}'::jsonb
                    );
                    EXIT;
                EXCEPTION WHEN unique_violation THEN
                    v_attempts := v_attempts + 1;
                    IF v_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) = 6 THEN
                invite_code_value := upper(invite_code_value);
                IF invite_code_value ~ '^[ABCDEFGHJKMNPQRSTUVWXYZ23456789]{6}$' 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 v_referrer_id;

                    IF v_referrer_id IS NOT NULL THEN
                        UPDATE public.profiles
                        SET referred_by = v_referrer_id
                        WHERE id = new.id;
                    END IF;
                END IF;
            END IF;

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

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


def downgrade() -> None:
    op.execute("DROP TRIGGER IF EXISTS on_auth_user_created ON auth.users")

    op.execute(
        """
        CREATE OR REPLACE FUNCTION public.initialize_profile_and_points_on_signup()
        RETURNS trigger
        LANGUAGE plpgsql
        SECURITY DEFINER
        SET search_path = public
        AS $$
        DECLARE
            v_username text;
            v_ledger_id uuid;
            v_event_id text;
        BEGIN
            v_username := 'user_' || substring(md5(new.id::text || clock_timestamp()::text || random()::text) from 1 for 6);
            v_ledger_id := md5(new.id::text || 'ledger' || clock_timestamp()::text || random()::text)::uuid;
            v_event_id := 'register:' || new.id::text;

            INSERT INTO public.profiles (id, username, avatar_url, bio, settings)
            VALUES (
                new.id,
                v_username,
                null,
                null,
                jsonb_build_object(
                    'version', 1,
                    'preferences', jsonb_build_object(
                        'interface_language', 'zh-CN',
                        'ai_language', 'zh-CN',
                        'timezone', 'Asia/Shanghai',
                        'country', 'CN'
                    ),
                    'privacy', jsonb_build_object('profile_visibility', 'public'),
                    'notification', jsonb_build_object('push_enabled', true)
                )
            )
            ON CONFLICT (id) DO NOTHING;

            INSERT INTO public.user_points (
                user_id,
                balance,
                frozen_balance,
                lifetime_earned,
                lifetime_spent,
                version
            )
            VALUES (new.id, 100, 0, 100, 0, 0)
            ON CONFLICT (user_id) DO NOTHING;

            INSERT INTO public.points_ledger (
                id,
                user_id,
                direction,
                amount,
                balance_after,
                change_type,
                biz_type,
                biz_id,
                event_id,
                operator_id,
                metadata
            )
            VALUES (
                v_ledger_id,
                new.id,
                1,
                100,
                100,
                'register',
                null,
                null,
                v_event_id,
                null,
                jsonb_build_object(
                    'schema_version', 1,
                    'reason_code', 'REGISTER_WELCOME',
                    'operator_type', 'system',
                    'run_id', v_event_id,
                    'request_id', null,
                    'ext', jsonb_build_object('source', 'auth_signup')
                )
            )
            ON CONFLICT (user_id, event_id) DO NOTHING;

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

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

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