"""给已经上线的库补上认证字段。create_all 不会修改旧表。"""

from sqlalchemy import inspect, text

from app.database import engine


def upgrade_schema() -> None:
    inspector = inspect(engine)
    tables = set(inspector.get_table_names())
    statements: list[str] = []
    need_username_index = False
    if "users" in tables:
        columns = {column["name"] for column in inspector.get_columns("users")}
        if "username" not in columns:
            statements.append("ALTER TABLE users ADD COLUMN username VARCHAR(32)")
        if "password_hash" not in columns:
            statements.append("ALTER TABLE users ADD COLUMN password_hash VARCHAR(255)")
        if "role" not in columns:
            statements.append("ALTER TABLE users ADD COLUMN role VARCHAR(16) DEFAULT 'user' NOT NULL")
        need_username_index = not _has_username_index(inspector)
    if "compute_nodes" in tables:
        node_columns = {column["name"] for column in inspector.get_columns("compute_nodes")}
        if "node_token" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN node_token VARCHAR(128)")
            statements.append(
                "CREATE UNIQUE INDEX IF NOT EXISTS uq_compute_nodes_node_token ON compute_nodes (node_token)"
            )
    if not statements and not need_username_index:
        return
    with engine.begin() as conn:
        for statement in statements:
            conn.execute(text(statement))
        if need_username_index:
            conn.execute(text("CREATE UNIQUE INDEX IF NOT EXISTS uq_users_username ON users (username)"))


def _has_username_index(inspector) -> bool:
    if "users" not in inspector.get_table_names():
        return False
    for index in inspector.get_indexes("users"):
        if index.get("column_names") == ["username"] and index.get("unique"):
            return True
    return False