"""Create user, chat, points, and invite schema.

Revision ID: 20260428_squash_0002
Revises: 20260428_squash_0001
Create Date: 2026-04-11 00:10:00

Squashed history: builds the final profiles settings shape directly, removes
the obsolete points_ledger.biz_id FK, and creates the final signup trigger.
"""

from typing import Sequence, Union

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

revision: str = "20260428_squash_0002"
down_revision: Union[str, Sequence[str], None] = "20260428_squash_0001"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    _create_profiles()
    _create_chat_tables()
    _create_points_tables()
    _create_invite_codes()
    _create_signup_helpers()


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

    for table_name in [
        "invite_codes",
        "register_bonus_claims",
        "points_audit_ledger",
        "points_ledger",
        "user_points",
        "messages",
        "sessions",
        "profiles",
    ]:
        _drop_service_only_rls(table_name)
    op.drop_table("invite_codes")
    op.drop_table("register_bonus_claims")
    op.drop_table("points_audit_ledger")
    op.drop_table("points_ledger")
    op.drop_table("user_points")
    op.drop_table("messages")
    op.drop_table("sessions")
    op.drop_table("profiles")


def _create_profiles() -> None:
    op.create_table(
        "profiles",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("username", sa.String(length=30), nullable=False),
        sa.Column("avatar_url", sa.Text(), nullable=True),
        sa.Column("bio", sa.String(length=200), nullable=True),
        sa.Column("settings", postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), nullable=False),
        sa.Column("referred_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.CheckConstraint("char_length(username) >= 1", name="ck_profiles_username_non_empty"),
        sa.ForeignKeyConstraint(["id"], ["auth.users.id"], ondelete="CASCADE"),
        sa.ForeignKeyConstraint(["referred_by"], ["profiles.id"], ondelete="SET NULL"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("ix_profiles_username", "profiles", ["username"])
    op.create_index("ix_profiles_settings_gin", "profiles", ["settings"], postgresql_using="gin")
    op.create_index("ix_profiles_referred_by", "profiles", ["referred_by"])
    _enable_service_only_rls("profiles")


def _create_chat_tables() -> 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(12, 6), server_default=sa.text("0"), nullable=False),
        sa.Column("state_snapshot", 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.CheckConstraint("session_type in ('chat', 'automation')", name="ck_sessions_session_type"),
        sa.CheckConstraint("status in ('pending', 'running', 'completed', 'failed')", name="ck_sessions_status"),
        sa.CheckConstraint("message_count >= 0", name="ck_sessions_message_count_non_negative"),
        sa.CheckConstraint("total_tokens >= 0", name="ck_sessions_total_tokens_non_negative"),
        sa.CheckConstraint("total_cost >= 0", name="ck_sessions_total_cost_non_negative"),
        sa.ForeignKeyConstraint(["user_id"], ["auth.users.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("ix_sessions_user_id", "sessions", ["user_id"])
    op.create_index("ix_sessions_user_activity", "sessions", ["user_id", "last_activity_at"])
    _enable_service_only_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(), server_default=sa.text("0"), nullable=False),
        sa.Column("output_tokens", sa.Integer(), server_default=sa.text("0"), nullable=False),
        sa.Column("cost", sa.Numeric(12, 6), server_default=sa.text("0"), nullable=False),
        sa.Column("latency_ms", sa.Integer(), nullable=True),
        sa.Column("visibility_mask", sa.BigInteger(), server_default=sa.text("0"), nullable=False),
        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.CheckConstraint("seq > 0", name="ck_messages_seq_positive"),
        sa.CheckConstraint("role in ('user', 'assistant', 'system', 'tool')", name="ck_messages_role"),
        sa.CheckConstraint("input_tokens >= 0", name="ck_messages_input_tokens_non_negative"),
        sa.CheckConstraint("output_tokens >= 0", name="ck_messages_output_tokens_non_negative"),
        sa.CheckConstraint("cost >= 0", name="ck_messages_cost_non_negative"),
        sa.CheckConstraint("latency_ms is null or latency_ms >= 0", name="ck_messages_latency_non_negative"),
        sa.ForeignKeyConstraint(["session_id"], ["sessions.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("session_id", "seq", name="uq_messages_session_seq"),
    )
    op.create_index("ix_messages_session_id", "messages", ["session_id"])
    op.create_index("ix_messages_session_seq_visibility", "messages", ["session_id", "seq", "visibility_mask"])
    _enable_service_only_rls("messages")


def _create_points_tables() -> None:
    op.create_table(
        "user_points",
        sa.Column("user_id", sa.UUID(), nullable=False),
        sa.Column("balance", sa.BigInteger(), server_default=sa.text("0"), nullable=False),
        sa.Column("frozen_balance", sa.BigInteger(), server_default=sa.text("0"), nullable=False),
        sa.Column("lifetime_earned", sa.BigInteger(), server_default=sa.text("0"), nullable=False),
        sa.Column("lifetime_spent", sa.BigInteger(), server_default=sa.text("0"), nullable=False),
        sa.Column("version", sa.Integer(), 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.CheckConstraint("balance >= 0", name="ck_user_points_balance_non_negative"),
        sa.CheckConstraint("frozen_balance >= 0", name="ck_user_points_frozen_balance_non_negative"),
        sa.CheckConstraint("lifetime_earned >= 0", name="ck_user_points_lifetime_earned_non_negative"),
        sa.CheckConstraint("lifetime_spent >= 0", name="ck_user_points_lifetime_spent_non_negative"),
        sa.CheckConstraint("frozen_balance <= balance", name="ck_user_points_frozen_le_balance"),
        sa.ForeignKeyConstraint(["user_id"], ["auth.users.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("user_id"),
    )
    _enable_service_only_rls("user_points")

    op.create_table(
        "points_ledger",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("user_id", sa.UUID(), nullable=False),
        sa.Column("direction", sa.SmallInteger(), nullable=False),
        sa.Column("amount", sa.BigInteger(), nullable=False),
        sa.Column("balance_after", sa.BigInteger(), nullable=False),
        sa.Column("change_type", sa.String(length=16), nullable=False),
        sa.Column("biz_type", sa.String(length=16), nullable=True),
        sa.Column("biz_id", sa.UUID(), nullable=True),
        sa.Column("event_id", sa.String(length=64), nullable=False),
        sa.Column("operator_id", sa.UUID(), nullable=True),
        sa.Column("metadata", postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), 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.CheckConstraint("amount > 0", name="ck_points_ledger_amount_positive"),
        sa.CheckConstraint("direction in (1, -1)", name="ck_points_ledger_direction_valid"),
        sa.CheckConstraint("balance_after >= 0", name="ck_points_ledger_balance_after_non_negative"),
        sa.CheckConstraint("change_type in ('register', 'consume', 'adjust', 'purchase', 'refund')", name="ck_points_ledger_change_type"),
        sa.CheckConstraint("biz_type is null or biz_type in ('chat', 'payment')", name="ck_points_ledger_biz_type"),
        sa.CheckConstraint("((change_type in ('register', 'adjust') and biz_type is null and biz_id is null) or (change_type = 'consume' and biz_type = 'chat' and biz_id is not null) or (change_type in ('purchase', 'refund') and biz_type = 'payment' and biz_id is not null))", name="ck_points_ledger_biz_binding"),
        sa.CheckConstraint("((change_type in ('register', 'purchase') and direction = 1) or (change_type in ('consume', 'refund') and direction = -1) or (change_type = 'adjust' and direction in (1, -1)))", name="ck_points_ledger_direction_by_change_type"),
        sa.CheckConstraint("jsonb_typeof(metadata) = 'object'", name="ck_points_ledger_metadata_object"),
        sa.CheckConstraint("metadata->>'schema_version' = '1' and metadata->>'operator_type' in ('user', 'system', 'admin') and coalesce(metadata->>'run_id', '') <> '' and (not (metadata ? 'ext') or jsonb_typeof(metadata->'ext') = 'object')", name="ck_points_ledger_metadata_common"),
        sa.CheckConstraint("(change_type <> 'register' or not (metadata ? 'charge'))", name="ck_points_ledger_metadata_register_shape"),
        sa.CheckConstraint("(change_type <> 'consume' or ((metadata ? 'charge') and jsonb_typeof(metadata->'charge') = 'object' and (metadata->'charge' ? 'message_id') and (metadata->'charge' ? 'message_seq') and (metadata->'charge' ? 'model_code') and (metadata->'charge' ? 'input_tokens') and (metadata->'charge' ? 'output_tokens') and (metadata->'charge' ? 'cost')))", name="ck_points_ledger_metadata_consume_shape"),
        sa.CheckConstraint("(change_type <> 'adjust' or ((metadata ? 'ext') and (metadata->'ext' ? 'reason') and coalesce(metadata #>> '{ext,reason}', '') <> ''))", name="ck_points_ledger_metadata_adjust_shape"),
        sa.CheckConstraint("(change_type not in ('purchase', 'refund') or ((metadata ? 'ext') and (metadata->'ext' ? 'source') and (metadata->'ext' ? 'platform') and (metadata->'ext' ? 'product_code') and (metadata->'ext' ? 'transaction_id') and coalesce(metadata #>> '{ext,source}', '') <> '' and coalesce(metadata #>> '{ext,platform}', '') <> '' and coalesce(metadata #>> '{ext,product_code}', '') <> '' and coalesce(metadata #>> '{ext,transaction_id}', '') <> ''))", name="ck_points_ledger_metadata_payment_shape"),
        sa.CheckConstraint("(change_type <> 'refund' or ((metadata ? 'ext') and (metadata->'ext' ? 'original_event_id') and coalesce(metadata #>> '{ext,original_event_id}', '') <> ''))", name="ck_points_ledger_metadata_refund_shape"),
        sa.ForeignKeyConstraint(["operator_id"], ["auth.users.id"], ondelete="SET NULL"),
        sa.ForeignKeyConstraint(["user_id"], ["auth.users.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("user_id", "event_id", name="uq_points_ledger_user_event"),
    )
    op.create_index("ix_points_ledger_user_created_at", "points_ledger", ["user_id", sa.text("created_at DESC")])
    op.create_index("ix_points_ledger_biz_type_biz_id", "points_ledger", ["biz_type", "biz_id"])
    _enable_service_only_rls("points_ledger")

    op.create_table(
        "points_audit_ledger",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("event_id", sa.String(length=64), nullable=False),
        sa.Column("user_id_snapshot", sa.UUID(), nullable=True),
        sa.Column("user_email_snapshot", sa.Text(), nullable=True),
        sa.Column("change_type", sa.String(length=16), nullable=False),
        sa.Column("biz_type", sa.String(length=16), nullable=True),
        sa.Column("biz_id", sa.UUID(), nullable=True),
        sa.Column("direction", sa.SmallInteger(), nullable=False),
        sa.Column("amount", sa.BigInteger(), nullable=False),
        sa.Column("balance_after", sa.BigInteger(), nullable=False),
        sa.Column("billed_to", sa.String(length=16), nullable=False),
        sa.Column("run_id", sa.String(length=128), nullable=True),
        sa.Column("request_id", sa.String(length=128), nullable=True),
        sa.Column("input_tokens", sa.Integer(), server_default=sa.text("0"), nullable=False),
        sa.Column("output_tokens", sa.Integer(), server_default=sa.text("0"), nullable=False),
        sa.Column("cost", sa.Numeric(12, 6), server_default=sa.text("0"), nullable=False),
        sa.Column("metadata", postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), 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.CheckConstraint("amount >= 0", name="ck_points_audit_ledger_amount_non_negative"),
        sa.CheckConstraint("direction in (1, 0, -1)", name="ck_points_audit_ledger_direction_valid"),
        sa.CheckConstraint("balance_after >= 0", name="ck_points_audit_ledger_balance_after_non_negative"),
        sa.CheckConstraint("change_type in ('register', 'consume', 'adjust', 'purchase', 'refund')", name="ck_points_audit_ledger_change_type"),
        sa.CheckConstraint("biz_type is null or biz_type in ('chat', 'payment')", name="ck_points_audit_ledger_biz_type"),
        sa.CheckConstraint("billed_to in ('user', 'platform')", name="ck_points_audit_ledger_billed_to"),
        sa.CheckConstraint("jsonb_typeof(metadata) = 'object'", name="ck_points_audit_ledger_metadata_object"),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("event_id", name="uq_points_audit_ledger_event_id"),
    )
    op.create_index("ix_points_audit_ledger_user_id_created_at", "points_audit_ledger", ["user_id_snapshot", sa.text("created_at DESC")])
    op.create_index("ix_points_audit_ledger_change_type_created_at", "points_audit_ledger", ["change_type", sa.text("created_at DESC")])
    _enable_service_only_rls("points_audit_ledger")

    op.create_table(
        "register_bonus_claims",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("email_hash", sa.String(length=64), nullable=False),
        sa.Column("user_email_snapshot", sa.Text(), nullable=False),
        sa.Column("first_user_id_snapshot", sa.UUID(), nullable=True),
        sa.Column("balance_snapshot", sa.BigInteger(), nullable=True),
        sa.Column("grant_event_id", sa.String(length=64), nullable=False),
        sa.Column("has_purchased_starter_pack", sa.Boolean(), server_default=sa.text("false"), 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"),
        sa.UniqueConstraint("email_hash", name="uq_register_bonus_claims_email_hash"),
        sa.UniqueConstraint("grant_event_id", name="uq_register_bonus_claims_grant_event_id"),
    )
    _enable_service_only_rls("register_bonus_claims")


def _create_invite_codes() -> 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'")
    _enable_service_only_rls("invite_codes")


def _create_signup_helpers() -> None:
    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_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);

            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('language', 'zh-CN', 'timezone', 'Asia/Shanghai'),
                    'privacy', jsonb_build_object('can_sell', false, 'profile_visibility', 'public'),
                    'notification', jsonb_build_object('allow_notifications', true, 'allow_vibration', true),
                    'divination_tutorial', jsonb_build_object(
                        'divination_entry_shown', false,
                        'auto_divination_shown', false,
                        'manual_divination_shown', false
                    )
                )
            )
            ON CONFLICT (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 _enable_service_only_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_service_only_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} DISABLE ROW LEVEL SECURITY")