import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;

/** Apply admin_ai_model.input_modalities migration to master DB. */
public class ApplyAdminAiModelModalities {
    private static final String HOST = "cq-cdb-8fjmemkb.sql.tencentcdb.com:27220";
    private static final String USER = "root";
    private static final String PASS = "Ylrz_1q2w3e4r5t6y";

    public static void main(String[] args) throws Exception {
        Class.forName("com.mysql.cj.jdbc.Driver");
        String db = args.length > 0 ? args[0] : "ylrz_saas";
        String url = "jdbc:mysql://" + HOST + "/" + db
                + "?useSSL=false&serverTimezone=GMT%2B8&characterEncoding=utf8";
        try (Connection c = DriverManager.getConnection(url, USER, PASS);
             Statement st = c.createStatement()) {
            if (!columnExists(st, db, "admin_ai_model", "input_modalities")) {
                st.executeUpdate("ALTER TABLE admin_ai_model ADD COLUMN input_modalities VARCHAR(100) "
                        + "DEFAULT 'text' COMMENT 'text,vision,audio comma-separated' AFTER status");
                System.out.println("[OK] added column input_modalities");
            } else {
                System.out.println("[SKIP] column input_modalities exists");
            }
            int v = st.executeUpdate("UPDATE admin_ai_model SET input_modalities='text,vision' WHERE ("
                    + "LOWER(model_identifier) LIKE '%vision%' OR LOWER(model_identifier) LIKE '%-vl%' "
                    + "OR LOWER(model_identifier) LIKE 'vl-%' OR LOWER(model_identifier) LIKE 'qwen-vl%' "
                    + "OR LOWER(model_identifier) LIKE 'hunyuan-vision%')");
            System.out.println("[OK] vision models updated: " + v);
            int t = st.executeUpdate("UPDATE admin_ai_model SET input_modalities='text' "
                    + "WHERE input_modalities IS NULL OR input_modalities=''");
            System.out.println("[OK] text default updated: " + t);
            System.out.println("Done " + db);
        }
    }

    private static boolean columnExists(Statement st, String db, String table, String column) throws Exception {
        ResultSet rs = st.executeQuery("SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA='"
                + db.replace("'", "''") + "' AND TABLE_NAME='" + table + "' AND COLUMN_NAME='" + column + "'");
        rs.next();
        return rs.getInt(1) > 0;
    }
}