2 files changed, 47 insertions, 1 deletions

diff --git a/modules/g2_import/helpers/g2_import.php b/modules/g2_import/helpers/g2_import.php
index 03c125ed..90f39ec1 100644
--- a/modules/g2_import/helpers/g2_import.php
+++ b/modules/g2_import/helpers/g2_import.php
@@ -271,6 +271,40 @@ class g2_import_Core {
     // @todo import album highlights
   }
 
+  /**
+   * Set the highlight properly for a single album
+   */
+  static function set_album_highlight(&$queue) {
+    // Dequeue the current album and enqueue its children
+    list($g2_album_id, $children) = each($queue);
+    unset($queue[$g2_album_id]);
+    foreach ($children as $key => $value) {
+      $queue[$key] = $value;
+    }
+
+    $g3_album_id = self::map($g2_album_id);
+    if (!$g3_album_id) {
+      return;
+    }
+
+    $table = g2(GalleryCoreApi::fetchThumbnailsByItemIds(array($g2_album_id)));
+    if (isset($table[$g2_album_id])) {
+      // Backtrack the source id to an item
+      $g2_source = $table[$g2_album_id];
+      while (GalleryUtilities::isA($g2_source, "GalleryDerivative")) {
+        $g2_source = g2(GalleryCoreApi::loadEntitiesById($g2_source->getDerivativeSourceId()));
+      }
+      $item_id = self::map($g2_source->getId());
+      if ($item_id) {
+        $item = ORM::factory("item", $item_id);
+        $g2_album = ORM::factory("item", $g3_album_id);
+        $g2_album->album_cover_item_id = $item->id;
+        $g2_album->thumb_dirty = 1;
+        $g2_album->save();
+        graphics::generate($g2_album);
+      }
+    }
+  }
 
   /**
    * Import a single photo or movie.
diff --git a/modules/g2_import/helpers/g2_import_task.php b/modules/g2_import/helpers/g2_import_task.php
index 161f0709..4cd95581 100644
--- a/modules/g2_import/helpers/g2_import_task.php
+++ b/modules/g2_import/helpers/g2_import_task.php
@@ -47,6 +47,7 @@ class g2_import_task_Core {
       $stats["items"] = $stats["photos"] + $stats["movies"];
       unset($stats["photos"]);
       unset($stats["movies"]);
+      $stats["highlights"] = $stats["albums"];
       $task->set("stats", $stats);
 
       $task->set("total", $total = array_sum(array_values($stats)));
@@ -68,7 +69,7 @@ class g2_import_task_Core {
       }
     }
 
-    $modes = array("groups", "users", "albums", "items", "comments", "tags", "done");
+    $modes = array("groups", "users", "albums", "items", "comments", "tags", "highlights", "done");
     while (!$task->done && microtime(true) - $start < 1.5) {
       if ($done[$modes[$mode]] == $stats[$modes[$mode]]) {
         // Nothing left to do for this mode.  Advance.
@@ -152,6 +153,17 @@ class g2_import_task_Core {
 
         break;
 
+      case "highlights":
+        if (empty($queue)) {
+          $task->set("queue", $queue = g2(GalleryCoreApi::fetchAlbumTree()));
+        }
+        g2_import::set_album_highlight($queue);
+        $task->status = t(
+          "Album highlights (%count of %total)",
+          array("count" => $done["tags"] + 1, "total" => $stats["albums"]));
+
+        break;
+
       case "done":
         $task->status = t("Import complete");
         $task->done = true;