"""数据处理运行表的显式检查与安装命令。"""

from __future__ import annotations

import argparse
from urllib.parse import urlsplit

from app.modules.data_process.store import DataProcessStore

REQUIRED_TASK_COLUMNS = (
    "generation_run_id",
    "results_confirmed",
    "workflow_step",
    "preview_status",
    "preview_progress",
    "preview_run_id",
    "preview_failure_reason",
    "preview_total_files",
    "preview_completed_files",
)


def _target_label(database_url: str) -> str:
    parsed = urlsplit(database_url)
    database = parsed.path.strip("/") or "(unknown)"
    return f"{parsed.hostname or '(unknown)'}:{parsed.port or 5432}/{database}"


def _schema_ready(store: DataProcessStore) -> bool:
    with store.connect() as conn:
        row = conn.execute(
            """
            SELECT COUNT(*) = %s AS ready
            FROM information_schema.columns
            WHERE table_schema=current_schema()
              AND table_name='data_process_tasks'
              AND column_name = ANY(%s)
            """,
            (len(REQUIRED_TASK_COLUMNS), list(REQUIRED_TASK_COLUMNS)),
        ).fetchone()
    return bool(row and row["ready"])


def main() -> int:
    parser = argparse.ArgumentParser(
        description="检查或显式安装数据处理运行表（不会由应用启动自动执行）"
    )
    action = parser.add_mutually_exclusive_group(required=True)
    action.add_argument("--check", action="store_true", help="只读检查迁移是否已安装")
    action.add_argument("--apply", action="store_true", help="执行 002 数据处理迁移")
    parser.add_argument(
        "--yes",
        action="store_true",
        help="确认允许修改 DATABASE_URL 指向的数据库；与 --apply 同时使用",
    )
    args = parser.parse_args()

    store = DataProcessStore()
    target = _target_label(store.database_url)
    if args.check:
        ready = _schema_ready(store)
        print(f"数据处理 schema：{'已安装' if ready else '未安装'}；目标：{target}")
        return 0 if ready else 1
    if not args.yes:
        parser.error("--apply 必须同时提供 --yes，确认修改目标数据库")

    print(f"正在安装数据处理 schema；目标：{target}")
    store.ensure_schema()
    if not _schema_ready(store):
        raise RuntimeError("迁移执行后仍未检测到 generation_run_id")
    print("数据处理 schema 安装完成")
    return 0


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