hw/app/cli.py

from __future__ import annotations

from collections.abc import Iterable
import os
from pathlib import Path
import subprocess
import sys

import click
from flask import Flask, current_app
from flask.cli import with_appcontext
from sqlalchemy import inspect

from app.extensions import db
from app.models import Account, Household, HouseholdMember, OptionItem
from app.services import demo_login_lines, seed_demo_accounts, seed_demo_households, seed_stress_households

DEFAULT_ADMIN_USERNAME = "admin"
DEFAULT_ADMIN_DISPLAY_NAME = "系统管理员"
DEFAULT_ADMIN_PASSWORD = "ChangeMe123!"

BASE_OPTION_ITEMS: tuple[dict[str, str | int | bool | None], ...] = (
    {"option_group": "relation_category", "option_code": "relative", "option_label": "亲戚", "parent_code": None, "sort_order": 10},
    {"option_group": "relation_category", "option_code": "friend", "option_label": "朋友", "parent_code": None, "sort_order": 20},
    {"option_group": "relation_category", "option_code": "colleague", "option_label": "同事", "parent_code": None, "sort_order": 30},
    {"option_group": "gift_method", "option_code": "cash", "option_label": "现金", "parent_code": None, "sort_order": 10},
    {"option_group": "gift_method", "option_code": "transfer", "option_label": "转账", "parent_code": None, "sort_order": 20},
    {"option_group": "gift_scene", "option_code": "pre_wedding", "option_label": "婚前", "parent_code": None, "sort_order": 10},
    {"option_group": "gift_scene", "option_code": "wedding_day", "option_label": "婚礼当天", "parent_code": None, "sort_order": 20},
    {"option_group": "gift_scene", "option_code": "post_wedding", "option_label": "婚后补记", "parent_code": None, "sort_order": 30},
)


def register_commands(app: Flask) -> None:
    app.cli.add_command(init_db_command)
    app.cli.add_command(seed_admin_command)
    app.cli.add_command(seed_options_command)
    app.cli.add_command(seed_all_command)
    app.cli.add_command(seed_demo_command)
    app.cli.add_command(seed_stress_command)
    app.cli.add_command(dev_bootstrap_command)


@click.command("init-db")
@with_appcontext
def init_db_command() -> None:
    db.create_all()
    click.echo("Database tables created.")


def ensure_seed_tables_exist() -> None:
    inspector = inspect(db.engine)
    required_tables = {Account.__tablename__, OptionItem.__tablename__}
    existing_tables = set(inspector.get_table_names())

    if required_tables.issubset(existing_tables):
        return

    missing_tables = ", ".join(sorted(required_tables - existing_tables))
    raise click.ClickException(
        f"Missing database tables: {missing_tables}. Run 'flask db upgrade' or 'flask init-db' first.",
    )


def ensure_demo_seed_tables_exist() -> None:
    inspector = inspect(db.engine)
    required_tables = {
        Account.__tablename__,
        OptionItem.__tablename__,
        Household.__tablename__,
        HouseholdMember.__tablename__,
    }
    existing_tables = set(inspector.get_table_names())

    if required_tables.issubset(existing_tables):
        return

    missing_tables = ", ".join(sorted(required_tables - existing_tables))
    raise click.ClickException(
        f"Missing database tables: {missing_tables}. Run 'flask db upgrade' or 'flask init-db' first.",
    )


@click.command("seed-admin")
@click.option("--username", default=DEFAULT_ADMIN_USERNAME, show_default=True)
@click.option("--display-name", default=DEFAULT_ADMIN_DISPLAY_NAME, show_default=True)
@click.option("--password", default=DEFAULT_ADMIN_PASSWORD, show_default=True)
@with_appcontext
def seed_admin_command(username: str, display_name: str, password: str) -> None:
    ensure_seed_tables_exist()

    existing_account = db.session.execute(
        db.select(Account).where(Account.username == username),
    ).scalar_one_or_none()
    if existing_account is not None:
        click.echo(f"Admin account '{username}' already exists.")
        return

    account = Account()
    account.username = username
    account.display_name = display_name
    account.set_password(password)
    account.role = "admin"
    account.status = "active"
    db.session.add(account)
    db.session.commit()
    click.echo(f"Admin account '{username}' created.")


@click.command("seed-options")
@with_appcontext
def seed_options_command() -> None:
    ensure_seed_tables_exist()

    created_count = seed_option_items(BASE_OPTION_ITEMS)
    db.session.commit()
    click.echo(f"Seeded {created_count} option items.")


@click.command("seed-all")
@click.option("--username", default=DEFAULT_ADMIN_USERNAME, show_default=True)
@click.option("--display-name", default=DEFAULT_ADMIN_DISPLAY_NAME, show_default=True)
@click.option("--password", default=DEFAULT_ADMIN_PASSWORD, show_default=True)
@with_appcontext
def seed_all_command(username: str, display_name: str, password: str) -> None:
    ensure_seed_tables_exist()

    created_options = seed_option_items(BASE_OPTION_ITEMS)

    existing_account = db.session.execute(
        db.select(Account).where(Account.username == username),
    ).scalar_one_or_none()
    if existing_account is None:
        admin_account = Account()
        admin_account.username = username
        admin_account.display_name = display_name
        admin_account.set_password(password)
        admin_account.role = "admin"
        admin_account.status = "active"
        db.session.add(admin_account)
        admin_created = True
    else:
        admin_created = False

    db.session.commit()

    if admin_created:
        click.echo(f"Admin account '{username}' created.")
    else:
        click.echo(f"Admin account '{username}' already exists.")
    click.echo(f"Seeded {created_options} option items.")


@click.command("seed-demo")
@with_appcontext
def seed_demo_command() -> None:
    ensure_demo_seed_tables_exist()

    created_accounts = seed_demo_accounts()
    created_households = seed_demo_households()
    db.session.commit()

    click.echo(f"Seeded {created_accounts} demo accounts.")
    click.echo(f"Seeded {created_households} demo households.")
    click.echo("Demo login accounts:")
    for login_line in demo_login_lines():
        click.echo(f"- {login_line}")


@click.command("seed-stress")
@click.option("--household-count", default=1000, show_default=True, type=int)
@click.option("--members-per-household", default=4, show_default=True, type=int)
@click.option("--prefix", default="STRESS-", show_default=True)
@with_appcontext
def seed_stress_command(household_count: int, members_per_household: int, prefix: str) -> None:
    ensure_demo_seed_tables_exist()

    try:
        created_households, created_members = seed_stress_households(
            household_count=household_count,
            members_per_household=members_per_household,
            household_code_prefix=prefix,
        )
    except ValueError as exc:
        raise click.ClickException(str(exc)) from exc

    db.session.commit()
    click.echo(f"Seeded {created_households} stress households.")
    click.echo(f"Seeded {created_members} stress household members.")


@click.command("dev-bootstrap")
@click.option("--host", default="127.0.0.1", show_default=True)
@click.option("--port", default=5000, show_default=True, type=int)
@click.option("--skip-server", is_flag=True, default=False)
@with_appcontext
def dev_bootstrap_command(host: str, port: int, skip_server: bool) -> None:
    app = current_app._get_current_object()
    project_root = os.fspath(Path(app.root_path).parent)
    python_executable = sys.executable
    env = os.environ.copy()
    env["FLASK_APP"] = "run.py"

    _run_python_command(
        [python_executable, "-m", "flask", "db", "upgrade"],
        env=env,
        cwd=project_root,
        label="Applying database migrations",
    )
    _run_python_command(
        [python_executable, "-m", "flask", "seed-all"],
        env=env,
        cwd=project_root,
        label="Seeding baseline admin and options",
    )
    _run_python_command(
        [python_executable, "-m", "flask", "seed-demo"],
        env=env,
        cwd=project_root,
        label="Seeding demo accounts and households",
    )

    click.echo(f"Service URL: http://{host}:{port}")
    click.echo("Available demo logins:")
    for login_line in demo_login_lines():
        click.echo(f"- {login_line}")

    if skip_server:
        click.echo("Skipped starting the Flask development server.")
        return

    click.echo("Starting Flask development server...")
    app.run(host=host, port=port, debug=True)


def seed_option_items(option_rows: Iterable[dict[str, str | int | bool | None]]) -> int:
    created_count = 0
    created_items_by_code: dict[tuple[str, str], OptionItem] = {}

    for row in option_rows:
        option_group = str(row["option_group"])
        option_code = str(row["option_code"])
        existing_item = db.session.execute(
            db.select(OptionItem).where(
                OptionItem.option_group == option_group,
                OptionItem.option_code == option_code,
            ),
        ).scalar_one_or_none()
        if existing_item is not None:
            created_items_by_code[(option_group, option_code)] = existing_item
            continue

        parent_id = None
        parent_code = row.get("parent_code")
        if parent_code is not None:
            parent_item = created_items_by_code.get((option_group.replace("_detail", "_category"), str(parent_code)))
            if parent_item is None:
                parent_item = db.session.execute(
                    db.select(OptionItem).where(OptionItem.option_code == str(parent_code)),
                ).scalar_one_or_none()
            if parent_item is not None:
                parent_id = parent_item.id

        option_item = OptionItem()
        option_item.option_group = option_group
        option_item.option_code = option_code
        option_item.option_label = str(row["option_label"])
        option_item.parent_id = parent_id
        option_item.sort_order = int(row.get("sort_order", 0) or 0)
        option_item.is_enabled = bool(row.get("is_enabled", True))
        option_item.is_system = bool(row.get("is_system", True))
        db.session.add(option_item)
        db.session.flush()
        created_items_by_code[(option_group, option_code)] = option_item
        created_count += 1

    return created_count


def _run_python_command(
    args: list[str],
    *,
    env: dict[str, str],
    cwd: str,
    label: str,
) -> None:
    click.echo(label)
    completed_process = subprocess.run(
        args,
        cwd=cwd,
        env=env,
        check=False,
        capture_output=True,
        text=True,
    )
    if completed_process.stdout:
        click.echo(completed_process.stdout.rstrip())
    if completed_process.returncode == 0:
        return

    if completed_process.stderr:
        click.echo(completed_process.stderr.rstrip(), err=True)
    raise click.ClickException(f"Command failed: {' '.join(args)}")