hw/app/services/demo_seed.py

from __future__ import annotations

from dataclasses import dataclass
from decimal import Decimal

from app.extensions import db
from app.models import Account, Household, HouseholdMember, OptionItem


DEMO_ADMIN_USERNAME = "admin"

DEFAULT_STRESS_HOUSEHOLD_COUNT = 1000
DEFAULT_STRESS_MEMBERS_PER_HOUSEHOLD = 4
DEFAULT_STRESS_HOUSEHOLD_CODE_PREFIX = "STRESS-"

_STRESS_SIDE_OPTIONS: tuple[str, ...] = ("groom_side", "bride_side", "both_side")
_STRESS_RELATION_CATEGORY_CODES: tuple[str, ...] = ("relative", "friend", "colleague")
_STRESS_TAG_CODE_SETS: tuple[tuple[str, ...], ...] = ((),)
_STRESS_INVITE_STATUS_OPTIONS: tuple[str, ...] = ("not_invited", "invited", "confirmed")
_STRESS_ATTENDANCE_STATUS_OPTIONS: tuple[str, ...] = ("pending", "attending", "partial", "absent")
_STRESS_GIVEN_STATUS_OPTIONS: tuple[str, ...] = ("not_given", "given")
_STRESS_CHILD_RED_PACKET_STATUS_OPTIONS: tuple[str, ...] = ("not_given", "partial", "given")
_STRESS_GIFT_METHOD_CODES: tuple[str, ...] = ("cash", "transfer")
_STRESS_GIFT_SCENE_CODES: tuple[str, ...] = ("pre_wedding", "wedding_day", "post_wedding")
_STRESS_RELATION_TO_HEAD_OPTIONS: tuple[str, ...] = (
    "本人",
    "配偶",
    "子女",
    "亲属",
    "同事",
    "同学",
)
_STRESS_GENDER_OPTIONS: tuple[str, ...] = ("male", "female")


@dataclass(frozen=True)
class DemoAccountSpec:
    username: str
    display_name: str
    password: str
    role: str


@dataclass(frozen=True)
class DemoMemberSpec:
    name: str
    relation_to_head: str | None
    gender: str | None
    age_group: str | None
    is_child: bool
    needs_red_packet: bool
    expected_to_attend: bool
    actually_attended: bool
    sort_order: int
    note: str | None


@dataclass(frozen=True)
class DemoHouseholdSpec:
    household_code: str
    head_name: str
    phone: str | None
    side: str
    relation_category_code: str | None
    relation_detail_code: str | None
    tag_codes: tuple[str, ...]
    invite_status: str
    attendance_status: str
    expected_attendee_count: int
    actual_attendee_count: int
    child_count: int
    red_packet_child_count: int
    total_gift_amount: Decimal
    gift_method_code: str | None
    gift_scene_code: str | None
    favor_status: str
    candy_status: str
    child_red_packet_status: str
    note: str | None
    members: tuple[DemoMemberSpec, ...]


DEMO_ACCOUNT_SPECS: tuple[DemoAccountSpec, ...] = (
    DemoAccountSpec(
        username="editor-demo",
        display_name="演示录入员",
        password="EditorDemo123!",
        role="editor",
    ),
    DemoAccountSpec(
        username="entry-demo",
        display_name="长辈演示账号",
        password="EntryDemo123!",
        role="entry_only",
    ),
    DemoAccountSpec(
        username="quick-editor-demo",
        display_name="快速录入账号",
        password="QuickEditor123!",
        role="quick_editor",
    ),
)


DEMO_HOUSEHOLD_SPECS: tuple[DemoHouseholdSpec, ...] = (
    DemoHouseholdSpec(
        household_code="D001",
        head_name="王大爷",
        phone="13800001234",
        side="groom_side",
        relation_category_code="relative",
        relation_detail_code=None,
        tag_codes=(),
        invite_status="confirmed",
        attendance_status="pending",
        expected_attendee_count=4,
        actual_attendee_count=0,
        child_count=1,
        red_packet_child_count=1,
        total_gift_amount=Decimal("666.00"),
        gift_method_code="cash",
        gift_scene_code="wedding_day",
        favor_status="not_given",
        candy_status="given",
        child_red_packet_status="not_given",
        note="演示数据：待电话确认当天到场时间。",
        members=(
            DemoMemberSpec(
                name="王大爷",
                relation_to_head="本人",
                gender="male",
                age_group="elder",
                is_child=False,
                needs_red_packet=False,
                expected_to_attend=True,
                actually_attended=False,
                sort_order=10,
                note="户主",
            ),
            DemoMemberSpec(
                name="王阿姨",
                relation_to_head="配偶",
                gender="female",
                age_group="elder",
                is_child=False,
                needs_red_packet=False,
                expected_to_attend=True,
                actually_attended=False,
                sort_order=20,
                note=None,
            ),
            DemoMemberSpec(
                name="王小宝",
                relation_to_head="孙辈",
                gender="male",
                age_group="child",
                is_child=True,
                needs_red_packet=True,
                expected_to_attend=True,
                actually_attended=False,
                sort_order=30,
                note="需准备儿童红包",
            ),
        ),
    ),
    DemoHouseholdSpec(
        household_code="D002",
        head_name="李阿姨",
        phone="13900004567",
        side="bride_side",
        relation_category_code="friend",
        relation_detail_code=None,
        tag_codes=(),
        invite_status="invited",
        attendance_status="attending",
        expected_attendee_count=2,
        actual_attendee_count=2,
        child_count=0,
        red_packet_child_count=0,
        total_gift_amount=Decimal("888.00"),
        gift_method_code="transfer",
        gift_scene_code="pre_wedding",
        favor_status="given",
        candy_status="given",
        child_red_packet_status="not_given",
        note="演示数据：已确认提前到场帮忙。",
        members=(
            DemoMemberSpec(
                name="李阿姨",
                relation_to_head="本人",
                gender="female",
                age_group="adult",
                is_child=False,
                needs_red_packet=False,
                expected_to_attend=True,
                actually_attended=True,
                sort_order=10,
                note=None,
            ),
            DemoMemberSpec(
                name="李叔叔",
                relation_to_head="配偶",
                gender="male",
                age_group="adult",
                is_child=False,
                needs_red_packet=False,
                expected_to_attend=True,
                actually_attended=True,
                sort_order=20,
                note=None,
            ),
        ),
    ),
    DemoHouseholdSpec(
        household_code="D003",
        head_name="陈老师",
        phone="13700007890",
        side="both_side",
        relation_category_code="colleague",
        relation_detail_code=None,
        tag_codes=(),
        invite_status="confirmed",
        attendance_status="partial",
        expected_attendee_count=3,
        actual_attendee_count=2,
        child_count=1,
        red_packet_child_count=0,
        total_gift_amount=Decimal("520.00"),
        gift_method_code="cash",
        gift_scene_code="wedding_day",
        favor_status="given",
        candy_status="not_given",
        child_red_packet_status="partial",
        note="演示数据：一位孩子不单独领红包。",
        members=(
            DemoMemberSpec(
                name="陈老师",
                relation_to_head="本人",
                gender="female",
                age_group="adult",
                is_child=False,
                needs_red_packet=False,
                expected_to_attend=True,
                actually_attended=True,
                sort_order=10,
                note=None,
            ),
            DemoMemberSpec(
                name="陈先生",
                relation_to_head="配偶",
                gender="male",
                age_group="adult",
                is_child=False,
                needs_red_packet=False,
                expected_to_attend=True,
                actually_attended=True,
                sort_order=20,
                note=None,
            ),
            DemoMemberSpec(
                name="陈小朋友",
                relation_to_head="子女",
                gender="female",
                age_group="child",
                is_child=True,
                needs_red_packet=False,
                expected_to_attend=True,
                actually_attended=False,
                sort_order=30,
                note="部分到场示例",
            ),
        ),
    ),
)


def seed_demo_accounts() -> int:
    created_count = 0
    for spec in DEMO_ACCOUNT_SPECS:
        existing_account = db.session.execute(
            db.select(Account).where(Account.username == spec.username),
        ).scalar_one_or_none()
        if existing_account is not None:
            continue

        account = Account(
            username=spec.username,
            display_name=spec.display_name,
            role=spec.role,
            status="active",
        )
        account.set_password(spec.password)
        db.session.add(account)
        created_count += 1

    return created_count


def seed_demo_households() -> int:
    created_count = 0
    option_index = _build_option_index()
    admin_account = _get_admin_account()
    admin_id = admin_account.id if admin_account is not None else None

    for spec in DEMO_HOUSEHOLD_SPECS:
        existing_household = db.session.execute(
            db.select(Household).where(Household.household_code == spec.household_code),
        ).scalar_one_or_none()
        if existing_household is not None:
            continue

        household = Household(
            household_code=spec.household_code,
            head_name=spec.head_name,
            phone=spec.phone,
            side=spec.side,
            relation_category_option_id=_resolve_option_id(option_index, "relation_category", spec.relation_category_code),
            relation_detail_option_id=_resolve_option_id(option_index, "relation_detail", spec.relation_detail_code),
            tag_option_ids_json=[
                option_id
                for option_id in (
                    _resolve_option_id(option_index, "tag", tag_code)
                    for tag_code in spec.tag_codes
                )
                if option_id is not None
            ],
            invite_status=spec.invite_status,
            attendance_status=spec.attendance_status,
            expected_attendee_count=spec.expected_attendee_count,
            actual_attendee_count=spec.actual_attendee_count,
            child_count=spec.child_count,
            red_packet_child_count=spec.red_packet_child_count,
            total_gift_amount=spec.total_gift_amount,
            gift_method_option_id=_resolve_option_id(option_index, "gift_method", spec.gift_method_code),
            gift_scene_option_id=_resolve_option_id(option_index, "gift_scene", spec.gift_scene_code),
            favor_status=spec.favor_status,
            candy_status=spec.candy_status,
            child_red_packet_status=spec.child_red_packet_status,
            note=spec.note,
            version=1,
            created_by=admin_id,
            updated_by=admin_id,
        )
        db.session.add(household)
        db.session.flush()

        for member_spec in spec.members:
            member = HouseholdMember(
                household_id=household.id,
                name=member_spec.name,
                relation_to_head=member_spec.relation_to_head,
                gender=member_spec.gender,
                age_group=member_spec.age_group,
                is_child=member_spec.is_child,
                needs_red_packet=member_spec.needs_red_packet,
                expected_to_attend=member_spec.expected_to_attend,
                actually_attended=member_spec.actually_attended,
                sort_order=member_spec.sort_order,
                note=member_spec.note,
                created_by=admin_id,
                updated_by=admin_id,
            )
            db.session.add(member)

        created_count += 1

    return created_count


def seed_stress_households(
    *,
    household_count: int = DEFAULT_STRESS_HOUSEHOLD_COUNT,
    members_per_household: int = DEFAULT_STRESS_MEMBERS_PER_HOUSEHOLD,
    household_code_prefix: str = DEFAULT_STRESS_HOUSEHOLD_CODE_PREFIX,
) -> tuple[int, int]:
    normalized_prefix = household_code_prefix.strip().upper()
    if household_count <= 0:
        raise ValueError("household_count must be greater than 0")
    if members_per_household <= 0:
        raise ValueError("members_per_household must be greater than 0")
    if not normalized_prefix:
        raise ValueError("household_code_prefix cannot be empty")

    created_household_count = 0
    created_member_count = 0
    option_index = _build_option_index()
    admin_account = _get_admin_account()
    admin_id = admin_account.id if admin_account is not None else None

    for index in range(1, household_count + 1):
        created = _create_stress_household(
            index=index,
            household_code=f"{normalized_prefix}{index:04d}",
            household_code_prefix=normalized_prefix,
            members_per_household=members_per_household,
            option_index=option_index,
            admin_id=admin_id,
        )
        if created is None:
            continue

        created_household_count += 1
        created_member_count += created

    return (created_household_count, created_member_count)


def _create_stress_household(
    *,
    index: int,
    household_code: str,
    household_code_prefix: str,
    members_per_household: int,
    option_index: dict[tuple[str, str], int],
    admin_id: int | None,
) -> int | None:
    existing_household = db.session.execute(
        db.select(Household).where(Household.household_code == household_code),
    ).scalar_one_or_none()
    if existing_household is not None:
        return None

    household = _build_stress_household(
        index=index,
        household_code=household_code,
        household_code_prefix=household_code_prefix,
        members_per_household=members_per_household,
        option_index=option_index,
        admin_id=admin_id,
    )
    db.session.add(household)
    db.session.flush()

    for member_index in range(1, members_per_household + 1):
        member = _build_stress_member(
            household_id=household.id,
            household_index=index,
            member_index=member_index,
            members_per_household=members_per_household,
            admin_id=admin_id,
        )
        db.session.add(member)

    return members_per_household


def _build_stress_household(
    *,
    index: int,
    household_code: str,
    household_code_prefix: str,
    members_per_household: int,
    option_index: dict[tuple[str, str], int],
    admin_id: int | None,
) -> Household:
    relation_category_code = _STRESS_RELATION_CATEGORY_CODES[(index - 1) % len(_STRESS_RELATION_CATEGORY_CODES)]
    side = _STRESS_SIDE_OPTIONS[(index - 1) % len(_STRESS_SIDE_OPTIONS)]
    invite_status = _STRESS_INVITE_STATUS_OPTIONS[(index - 1) % len(_STRESS_INVITE_STATUS_OPTIONS)]
    attendance_status = _STRESS_ATTENDANCE_STATUS_OPTIONS[(index - 1) % len(_STRESS_ATTENDANCE_STATUS_OPTIONS)]
    tag_codes = _STRESS_TAG_CODE_SETS[(index - 1) % len(_STRESS_TAG_CODE_SETS)]
    child_count = 1 if members_per_household >= 3 and index % 3 == 0 else 0
    red_packet_child_count = child_count if index % 2 == 0 else 0
    actual_attendee_count = _stress_actual_attendee_count(
        expected_attendee_count=members_per_household,
        attendance_status=attendance_status,
    )
    total_gift_amount = Decimal(index * 50).quantize(Decimal("0.01"))
    phone_number = f"139{index:08d}"[-11:]
    normalized_phone = f"1{phone_number[-10:]}"
    head_name_prefix = "" if household_code_prefix == DEFAULT_STRESS_HOUSEHOLD_CODE_PREFIX else household_code_prefix
    return Household(
        household_code=household_code,
        head_name=f"{head_name_prefix}压力户主{index:04d}",
        phone=normalized_phone,
        side=side,
        relation_category_option_id=_resolve_option_id(option_index, "relation_category", relation_category_code),
        relation_detail_option_id=None,
        tag_option_ids_json=[
            option_id
            for option_id in (
                _resolve_option_id(option_index, "tag", tag_code)
                for tag_code in tag_codes
            )
            if option_id is not None
        ],
        invite_status=invite_status,
        attendance_status=attendance_status,
        expected_attendee_count=members_per_household,
        actual_attendee_count=actual_attendee_count,
        child_count=child_count,
        red_packet_child_count=red_packet_child_count,
        total_gift_amount=total_gift_amount,
        gift_method_option_id=_resolve_option_id(
            option_index,
            "gift_method",
            _STRESS_GIFT_METHOD_CODES[(index - 1) % len(_STRESS_GIFT_METHOD_CODES)],
        ),
        gift_scene_option_id=_resolve_option_id(
            option_index,
            "gift_scene",
            _STRESS_GIFT_SCENE_CODES[(index - 1) % len(_STRESS_GIFT_SCENE_CODES)],
        ),
        favor_status=_STRESS_GIVEN_STATUS_OPTIONS[index % len(_STRESS_GIVEN_STATUS_OPTIONS)],
        candy_status=_STRESS_GIVEN_STATUS_OPTIONS[(index + 1) % len(_STRESS_GIVEN_STATUS_OPTIONS)],
        child_red_packet_status=_STRESS_CHILD_RED_PACKET_STATUS_OPTIONS[
            (index - 1) % len(_STRESS_CHILD_RED_PACKET_STATUS_OPTIONS)
        ],
        note=f"压力测试数据 #{index:04d}",
        version=1,
        created_by=admin_id,
        updated_by=admin_id,
    )


def _build_stress_member(
    *,
    household_id: int,
    household_index: int,
    member_index: int,
    members_per_household: int,
    admin_id: int | None,
) -> HouseholdMember:
    is_head = member_index == 1
    is_child = members_per_household >= 3 and member_index == members_per_household and household_index % 3 == 0
    attended = household_index % 4 in {1, 2} and not is_child
    return HouseholdMember(
        household_id=household_id,
        name=f"压力成员{household_index:04d}-{member_index}",
        relation_to_head=_stress_relation_to_head(member_index=member_index),
        gender=_STRESS_GENDER_OPTIONS[(household_index + member_index - 2) % len(_STRESS_GENDER_OPTIONS)],
        age_group="child" if is_child else _stress_age_group(member_index=member_index, is_head=is_head),
        is_child=is_child,
        needs_red_packet=is_child and household_index % 2 == 0,
        expected_to_attend=True,
        actually_attended=attended,
        sort_order=member_index * 10,
        note="压力测试儿童成员" if is_child else ("户主" if is_head else None),
        created_by=admin_id,
        updated_by=admin_id,
    )


def _stress_actual_attendee_count(*, expected_attendee_count: int, attendance_status: str) -> int:
    if attendance_status == "attending":
        return expected_attendee_count
    if attendance_status == "partial":
        return max(expected_attendee_count - 1, 1)
    if attendance_status == "absent":
        return 0
    return 0


def _stress_relation_to_head(*, member_index: int) -> str:
    if member_index == 1:
        return "本人"
    if member_index == 2:
        return "配偶"
    if member_index == 3:
        return "子女"
    return _STRESS_RELATION_TO_HEAD_OPTIONS[(member_index - 1) % len(_STRESS_RELATION_TO_HEAD_OPTIONS)]


def _stress_age_group(*, member_index: int, is_head: bool) -> str:
    if is_head:
        return "adult"
    if member_index == 2:
        return "adult"
    return "elder" if member_index % 2 == 0 else "adult"


def demo_login_lines() -> tuple[str, ...]:
    return (
        "admin / ChangeMe123!（管理员）",
        "editor-demo / EditorDemo123!（演示录入员）",
        "entry-demo / EntryDemo123!（长辈录入演示）",
        "quick-editor-demo / QuickEditor123!（快速录入账号）",
    )


def _build_option_index() -> dict[tuple[str, str], int]:
    option_items = db.session.execute(db.select(OptionItem)).scalars().all()
    return {
        (option_item.option_group, option_item.option_code): option_item.id
        for option_item in option_items
    }


def _resolve_option_id(
    option_index: dict[tuple[str, str], int],
    option_group: str,
    option_code: str | None,
) -> int | None:
    if option_code is None:
        return None

    return option_index.get((option_group, option_code))


def _get_admin_account() -> Account | None:
    return db.session.execute(
        db.select(Account).where(Account.username == DEMO_ADMIN_USERNAME),
    ).scalar_one_or_none()