clear_image_config.py

#!/usr/bin/env python3
"""Clear LandPPT image generation configuration from the local SQLite database.

Usage:
    python clear_image_config.py
    python clear_image_config.py --user-id 1
    python clear_image_config.py --include-system
    python clear_image_config.py --dry-run
"""

from __future__ import annotations

import argparse
import sqlite3
from pathlib import Path
from typing import Iterable

IMAGE_CONFIG_KEYS = (
    "gemini_image_api_base",
    "gemini_image_model",
    "openai_image_api_base",
    "openai_image_model",
    "openai_image_quality",
    "openai_image_size",
    "openai_api_key_image",
    "openai_image_api_key",
    "siliconflow_api_key",
    "pollinations_api_key",
    "gemini_image_api_key",
)

SECRET_KEYS = {
    "openai_api_key_image",
    "openai_image_api_key",
    "siliconflow_api_key",
    "pollinations_api_key",
    "gemini_image_api_key",
}


def mask_value(key: str, value: object) -> str:
    if value is None or str(value) == "":
        return "<EMPTY>"

    text = str(value)
    if key in SECRET_KEYS:
        if len(text) <= 10:
            return "***"
        return f"{text[:6]}...{text[-4:]}"

    return text


def format_user_id(user_id: object) -> str:
    return "SYSTEM" if user_id is None else str(user_id)


def placeholders(values: Iterable[object]) -> str:
    return ",".join("?" for _ in values)


def fetch_rows(conn: sqlite3.Connection, include_system: bool, user_id: int) -> list[tuple]:
    keys_sql = placeholders(IMAGE_CONFIG_KEYS)
    params: list[object] = list(IMAGE_CONFIG_KEYS)

    if include_system:
        scope_sql = "(user_id = ? OR user_id IS NULL)"
        params.append(user_id)
    else:
        scope_sql = "user_id = ?"
        params.append(user_id)

    sql = f"""
        SELECT user_id, config_key, config_value, config_type, category
        FROM user_configs
        WHERE category = 'image_service'
          AND config_key IN ({keys_sql})
          AND {scope_sql}
        ORDER BY user_id, config_key
    """
    return conn.execute(sql, params).fetchall()


def print_rows(title: str, rows: list[tuple]) -> None:
    print(f"\n{title}")
    print("-" * len(title))

    if not rows:
        print("No matching rows.")
        return

    for user_id, key, value, config_type, category in rows:
        print(
            f"user={format_user_id(user_id):>6} | "
            f"key={key:<24} | "
            f"value={mask_value(key, value)} | "
            f"type={config_type} | category={category}"
        )


def clear_config(conn: sqlite3.Connection, include_system: bool, user_id: int) -> int:
    keys_sql = placeholders(IMAGE_CONFIG_KEYS)
    params: list[object] = list(IMAGE_CONFIG_KEYS)

    if include_system:
        scope_sql = "(user_id = ? OR user_id IS NULL)"
        params.append(user_id)
    else:
        scope_sql = "user_id = ?"
        params.append(user_id)

    sql = f"""
        DELETE FROM user_configs
        WHERE category = 'image_service'
          AND config_key IN ({keys_sql})
          AND {scope_sql}
    """
    cursor = conn.execute(sql, params)
    return cursor.rowcount


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Clear LandPPT image generation configuration from landppt.db."
    )
    parser.add_argument(
        "--db",
        default="landppt.db",
        help="Path to SQLite database. Default: landppt.db",
    )
    parser.add_argument(
        "--user-id",
        type=int,
        default=1,
        help="User ID whose image config should be cleared. Default: 1",
    )
    parser.add_argument(
        "--include-system",
        action="store_true",
        help="Also delete system default image config where user_id is NULL.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Only show what would be deleted, without changing the database.",
    )
    args = parser.parse_args()

    db_path = Path(args.db)
    if not db_path.exists():
        print(f"Database not found: {db_path}")
        return 1

    conn = sqlite3.connect(str(db_path))
    try:
        before_rows = fetch_rows(conn, include_system=args.include_system, user_id=args.user_id)
        print_rows("Matching image config before cleanup", before_rows)

        if args.dry_run:
            print("\nDry run enabled. No rows were deleted.")
            return 0

        deleted_count = clear_config(conn, include_system=args.include_system, user_id=args.user_id)
        conn.commit()
        print(f"\nDeleted rows: {deleted_count}")

        after_rows = fetch_rows(conn, include_system=args.include_system, user_id=args.user_id)
        print_rows("Matching image config after cleanup", after_rows)

        print("\nDone. Restart LandPPT or refresh the settings page to see the result.")
        return 0
    finally:
        conn.close()


if __name__ == "__main__":
    raise SystemExit(main())